import pandas as pd
from datetime import datetime
import numpy

class AssetClassMomentumAlgorithm(QCAlgorithm):
    
    
    def Initialize(self):

        self.SetStartDate(2012, 1, 1)  
        self.SetEndDate(datetime.now()) 
        self.InitCash = 10000
        self.SetCash(self.InitCash)  
        self.MKT = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.mkt = [] # for plotting stretegy vs SPY benchmark
        
        # the only risk off ticker to park all money during risk off regime
        self.risk_off_ticker = self.AddEquity("TYD", Resolution.Hour).Symbol
        
        # brokerage model
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) 
        
        self.data_mom_resultant = {}
        self.bil_mom_resultant = 0
        self.dbb_mom_resultant = 0
        
                
        # UNIVERSE FOR RISK ON MARKET REGIME
        self.symbols_risk_on = ["XLF", # finance
                                "QQQ", # Nasdaq-100
                                "IJK", # mid-cap 400 growth
                                "IJS", # small-cap 600 value
                                "IJT", # small-cap 600 growth
                                "IJJ", # mid-cap 400 value
                                "SPYG", # US large-cap growth
                                "SPYV", # US large-cap value
                                "XHB" # homebuilders
                                "XRT", # retail
                                "SMH", # semiconductors
                                "PPH" # pharma
                                ]
        
        
        # Tickers used as indicators for Risk-On Risk-Off regimes
        self.risk_off_bil = self.AddEquity("BIL", Resolution.Hour).Symbol # cash equivalent
        self.risk_off_dbb = self.AddEquity("DBB", Resolution.Hour).Symbol # base metals (risk on regime)
        
        # Subscribe to all risk on tickers
        
        for symbol in self.symbols_risk_on:
            
            #subscribe to a symbol
            self.AddEquity(symbol, Resolution.Minute).Symbol
        
                        
        # warm up for 4 months each 21 business day
        self.SetWarmUp(int(4*21))
            
                
        # schedule for rebalance
        self.Schedule.On(self.DateRules.MonthEnd("SPY"), self.TimeRules.AfterMarketOpen("SPY", 10), 
            self.Rebalance)
            

    def Rebalance(self):
        
        if self.IsWarmingUp: return
        
               
        # CALC MOMENTUM SCORES
        
        
        for symbol in self.symbols_risk_on:
            
                
            # if tradable
            if self.Securities[symbol].IsTradable == True:
                
                # fetch history data only for a given ticker
                self.df = self.History([symbol], timedelta(days=int(40*21)), Resolution.Daily) # fetch data for the last 63 days for a given ticker
                #self.Debug(str(self.Time) + str(self.df))
                if len(self.df.index)<62:continue
                if not self.df.empty:
                    symbol_quotebar = self.df.loc[symbol]
                    
                    mom_one_month = (symbol_quotebar["close"][0] - symbol_quotebar["close"][20])/symbol_quotebar["close"][20] # 1st and 21st close price
                    mom_second_month = (symbol_quotebar["close"][21] - symbol_quotebar["close"][41])/symbol_quotebar["close"][41] # 22nd and 42nd price
                    mom_three_month = (symbol_quotebar["close"][42] - symbol_quotebar["close"][62])/symbol_quotebar["close"][62] # 43rd and 63rd price
                    
                    # calculate resultant MOMP and populate the dictionary
                    resultant_momentum = 3*mom_one_month + 2*mom_second_month + 1*mom_three_month
                
                    self.data_mom_resultant[symbol] = resultant_momentum
        
           
        # current market indicators risk on/off momentums and calculating the resultant momentums
        
        # BIL
        # fetch history data last 4 months
        self.df = self.History(self.risk_off_bil, timedelta(days=int(4*21)), Resolution.Daily)
        if len(self.df.index)<62:return
        # calc resultant mon for BIL
        symbol_quotebar = self.df.loc["BIL"]
                    
        mom_one_month = (symbol_quotebar["close"][0] - symbol_quotebar["close"][20])/symbol_quotebar["close"][20] # 1st and 21st close price
        mom_second_month = (symbol_quotebar["close"][21] - symbol_quotebar["close"][41])/symbol_quotebar["close"][41] # 22nd and 42nd price
        mom_three_month = (symbol_quotebar["close"][42] - symbol_quotebar["close"][62])/symbol_quotebar["close"][62] # 43rd and 63rd price
                    
        # calculate resultant MOMP
        bil_resultant_momentum = 3*mom_one_month + 2*mom_second_month + 1*mom_three_month
                
        self.bil_mom_resultant = bil_resultant_momentum
        
        # DBB
        # fetch history data last 4 months
        self.df = self.History("DBB", timedelta(days=int(4*21)), Resolution.Daily)
        if len(self.df.index)<62:return
        # calc resultant mon for DBB
        symbol_quotebar = self.df.loc["DBB"]
                    
        mom_one_month = (symbol_quotebar["close"][0] - symbol_quotebar["close"][20])/symbol_quotebar["close"][20] # 1st and 21st close price
        mom_second_month = (symbol_quotebar["close"][21] - symbol_quotebar["close"][41])/symbol_quotebar["close"][41] # 22nd and 42nd price
        mom_three_month = (symbol_quotebar["close"][42] - symbol_quotebar["close"][62])/symbol_quotebar["close"][62] # 43rd and 63rd price
                    
        # calculate resultant MOMP
        dbb_resultant_momentum = 3*mom_one_month + 2*mom_second_month + 1*mom_three_month
                
        self.dbb_mom_resultant = dbb_resultant_momentum
        
        
        
        
        # picking top performers by momentum by creating Series with top N rows in each Series for top performers, symbol is index, 1 column with mom
        top_symbols_risk_on = pd.Series(self.data_mom_resultant).sort_values(ascending = False)[:4] # returns Series of top momentums (a series with N rows), here we take only symbols of top N performers
        
        
        
        # Case 1: mom DBB > mom BIL (strategy is in risk on regime)
        if self.dbb_mom_resultant > self.bil_mom_resultant:
            
            # if Portfolio has open positions
            if self.Portfolio.Invested:
                
                for sec in self.Portfolio.Keys:
                    
                    if sec not in top_symbols_risk_on.index.tolist():   # here and throughout try to flatten index eg pd.Series.index.tolist()             
                        
                        self.Liquidate(sec)
                        
                        for sec in top_symbols_risk_on.index:        
                            self.SetHoldings(sec, 1 / len(top_symbols_risk_on))
                        
                    # if nothing to rebalance, i.e. open position(s) are same as suggested, continue keeping a sec as open position
                    else:
                        pass
                    
            # Portfolio is all cash, not invested
            else:
                # open positions from top n ETFs during risk on regime
                for sec in top_symbols_risk_on.index:        
                    self.SetHoldings(sec, 1 / len(top_symbols_risk_on))
                    
                
        # Case 2: mom DBB < mom BIL (strategy is in risk off regime)
        if self.dbb_mom_resultant <= self.bil_mom_resultant:
        
            # if Portfolio has open positions
            if self.Portfolio.Invested:
                
                for sec in self.Portfolio.Keys:
                    
                    # all positions are closed
                    self.Liquidate(sec)
                    
                # buy risk off ticker using all money
                self.SetHoldings(self.risk_off_ticker, 1)
                    
                    
                    
            # Portfolio is all cash, not invested
            else:
                # buy risk off ticker using all money
                self.SetHoldings(self.risk_off_ticker, 1)
                
# just plotting strategy vs the SPY benchmark
    def OnEndOfDay(self): 
        
        mkt_price = self.Securities[self.MKT].Close
        self.mkt.append(mkt_price)
        mkt_perf = self.InitCash * self.mkt[-1] / self.mkt[0] 
        self.Plot('Strategy Equity', self.MKT, mkt_perf)     
        
        account_leverage = self.Portfolio.TotalHoldingsValue / self.Portfolio.TotalPortfolioValue
        
        self.Plot('Holdings', 'leverage', round(account_leverage, 2))
        self.Plot('Holdings', 'Target Leverage', 1)