class SectorAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)  
        self.SetCash(10000) 
        self.Settings.FreePortfolioValuePercentage = 0.05
        
        # Daily indicator data.
        self.s_indicators = {}
        self.b_indicators = {}
        
        self.period = 200
        
        self.sectors_in = 0
        
        self.SetWarmUp(self.period)
        
        self.sectors = [
                        "XRLE", # Real Estate SPDR Fund
                        "XLK",  # Technology Select Sector SPDR Fund
                        "XLE",  # Energy Select Sector SPDR Fund
                        "XLV",  # Health Care Select Sector SPDR Fund
                        "XLF",  # Financial Select Sector SPDR Fund
                        "XLI",  # Industrials Select Sector SPDR Fund
                        "XLB",  # Materials Select Sector SPDR Fund
                        "XLY",  # Consumer Discretionary Select Sector SPDR Fund
                        "XLP",  # Consumer Staples Select Sector SPDR Fund
                        "XLU",  # Utilities Select Sector SPDR Fund
                        #"XLC",  # Communication Services SPDR Fund
                        ]  
                        
        self.bonds = ["TLT", # iShares 20 Plus Year Treasury Bond ETF
                    ]
        
        self.SetBenchmark('SPY')
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
                    
        for sector in self.sectors:
            symbol = self.AddEquity(sector, Resolution.Daily).Symbol
            self.s_indicators[symbol] = SymbolData(self, symbol, self.period)
            
        for bond in self.bonds:
            symbol = self.AddEquity(bond, Resolution.Daily).Symbol
            self.b_indicators[symbol] = SymbolData(self, symbol, self.period)
            
        self.plot_allocation = Chart('Allocation')
        
        self.Schedule.On(self.DateRules.MonthStart(self.sectors[0]), self.TimeRules.AfterMarketOpen(self.sectors[0]), self.Rebalance)
        #self.Schedule.On(self.DateRules.WeekStart(self.sectors[0]), self.TimeRules.AfterMarketOpen(self.sectors[0]), self.Rebalance)
        self.Schedule.On(self.DateRules.EveryDay(self.sectors[0]), self.TimeRules.BeforeMarketClose(self.sectors[0], 1), self.PlotData)
        

    def Rebalance(self):
        n_sectors = len(self.sectors)
        n_bonds = len(self.bonds)
        in_out_sectors = 0
        
        # Sectors ETF
        for sector in self.s_indicators:
            if self.s_indicators[sector].sell:
                if self.Portfolio[sector].Invested:
                    in_out_sectors -= 1
                    # Sell sector ETF
                    self.Liquidate(sector)
            if self.s_indicators[sector].buy:
                if not self.Portfolio[sector].Invested:
                    in_out_sectors += 1
                    # Buy sector ETF
                    self.SetHoldings(sector, 1 / n_sectors)
                        
        # Update number of long ETF
        self.sectors_in = self.sectors_in + in_out_sectors    
        
        # Bonds
        for bond in self.b_indicators:
            # Sell bonds
            if self.b_indicators[bond].sell and self.Portfolio[bond].Invested:
                # Sell bonds
                    self.Liquidate(bond)
            
            # if bond pct must change, adjust the size of position
            if in_out_sectors:
                if self.b_indicators[bond].buy:
                    # New size
                    bond_pct = ((n_sectors - self.sectors_in) / n_sectors) / n_bonds
                    self.SetHoldings(bond, bond_pct)
                                    
            
                        

    def PlotData(self):
        cash = self.Portfolio.Cash + self.Portfolio.UnsettledCash 
        
        sectors_allocation = 0
        for sector in self.s_indicators:
            sectors_allocation += self.Portfolio[sector].Price * self.Portfolio[sector].Quantity
            
        bonds_allocation = 0
        for bond in self.b_indicators:
            bonds_allocation += self.Portfolio[sector].Price * self.Portfolio[sector].Quantity
            
        self.Plot("Allocation", "Cash", cash/self.Portfolio.TotalPortfolioValue * 100)
        self.Plot("Allocation", "Sectors", sectors_allocation/self.Portfolio.TotalPortfolioValue * 100)
        self.Plot("Allocation", "Bonds", bonds_allocation/self.Portfolio.TotalPortfolioValue * 100)
        
        self.Plot("Sectors", "IN", self.sectors_in)
        

class SymbolData(object):
    def __init__(self, algo, symbol, period):
        self.algo = algo
        self.symbol = symbol
        self.slow_ema = ExponentialMovingAverage(period)
        self.fast_ema = ExponentialMovingAverage(int(period/10))
        algo.RegisterIndicator(symbol.Value, self.slow_ema, Resolution.Daily)
        algo.RegisterIndicator(symbol.Value, self.fast_ema, Resolution.Daily)
        
    @property
    def buy2(self):
        if self.slow_ema.IsReady:
            price = self.algo.Securities[self.symbol].Close
            indicator = self.slow_ema.Current.Value
            
            if price > indicator:
                return True
            else:
                return False
        else:
            None
                
    @property
    def buy(self):
        if self.slow_ema.IsReady and self.fast_ema.IsReady:
            slow = self.slow_ema.Current.Value
            #fast = self.fast_ema.Current.Value
            fast = self.algo.Securities[self.symbol].Close
            
            if fast > slow:
                return True
            else:
                return False
        else:
            None
                
    @property
    def sell(self):
        return not self.buy