#region imports
from AlgorithmImports import *
#endregion

#See:  https://papers.ssrn.com/sol3/papers.cfm?abstract_id=4166845
import pandas as pd
import numpy as np

class BoldAssetAllocation(QCAlgorithm):

    def Initialize(self):

        # Setting starting cap to 100000, which will be used in the SPY benchmark chart
        self.cap = 100000

        self.SetCash(self.cap)

        self.SetStartDate(2015,1,1)  # Set Start Date
        #self.SetStartDate(2022,1,1)  
        #self.SetStartDate(2023,1,1)  
        self.SetEndDate(2023, 7, 21)

        

        self.start_cash = self.cap
        self.SetCash(self.start_cash)  # Set Strategy Cash
        self.SetBenchmark('SPY')
        
        # Algo Parameters
        self.prds = [1,3,6,12]
        self.prdwts = np.array([12,6,2,1])
        
        self.LO, self.LD, self.LP, self.B, self.TO, self.TD = [12,12,0,1,2,3]
        self.hprd = max(self.prds+[self.LO,self.LD])*21+50
        
        # Assets
        #self.canary = ['SPY','EFA','EEM','BND']
        self.canary = ['TIP', 'EEM']

        #self.offensive = ['QQQ','EFA','EEM','BND', 'XLK', 'IWF']
        self.offensive = ['QQQ', 'XLK', 'MTUM']
        self.defensive = ['BIL','BND','DBC','IEF','LQD','TIP','TLT', 'UUP']
        #self.defensive = ['TLT', 'UUP']
        self.safe = 'BIL'
        # repeat safe asset so it can be selected multiple times
        self.alldefensive = self.defensive + [self.safe] * max(0,self.TD - sum([1*(e==self.safe) for e in self.defensive]))
        self.eqs = list(dict.fromkeys(self.canary+self.offensive+self.alldefensive))
        for eq in self.eqs:
            self.AddEquity(eq,Resolution.Minute)
        
        # Plot SPY on Equity Graph
        self.BNC = self.AddEquity('SPY',Resolution.Daily).Symbol
        self.mkt = []
        
        # monthly rebalance
        self.Schedule.On(self.DateRules.MonthStart(self.canary[0]),self.TimeRules.AfterMarketOpen(self.canary[0],30),self.rebal)
        self.Trade = True
        
    def rebal(self):
        self.Trade = True

    def OnData(self, data):           
        if self.Trade:
            # Get price data and trading weights
            h = self.History(self.eqs,self.hprd,Resolution.Daily)['close'].unstack(level=0)
            wts = self.trade_wts(h)

            # trade
            port_tgt = [PortfolioTarget(x,y) for x,y in zip(wts.index,wts.values)]
            self.SetHoldings(port_tgt)
            
            self.Trade = False
    
    def trade_wts(self,hist):
        # initialize wts Series
        wts = pd.Series(0,index=hist.columns)
        # end of month values
        h_eom = (hist.loc[hist.groupby(hist.index.to_period('M')).apply(lambda x: x.index.max())]
                .iloc[:-1,:])

        # =====================================
        # check if canary universe is triggered
        # =====================================
        # build dataframe of momentum values
        mom = h_eom.iloc[-1,:].div(h_eom.iloc[[-p-1 for p in self.prds],:],axis=0)-1
        mom = mom.loc[:,self.canary].T
        # Determine number of canary securities with negative weighted momentum
        n_canary = np.sum(np.sum(mom.values*self.prdwts,axis=1)<0)
        # % equity offensive 
        pct_in = 1-min(1,n_canary/self.B)

        # =====================================
        # get weights for offensive and defensive universes
        # =====================================
        # determine weights of offensive universe
        if pct_in > 0:
            # price / SMA
            mom_in = h_eom.iloc[-1,:].div(h_eom.iloc[[-t for t in range(1,self.LO+1)]].mean(axis=0),axis=0)
            mom_in = mom_in.loc[self.offensive].sort_values(ascending=False)
            # equal weightings to top relative momentum securities
            in_wts = pd.Series(pct_in/self.TO,index=mom_in.index[:self.TO])
            wts = pd.concat([wts,in_wts])
        # determine weights of defensive universe
        if pct_in < 1:
            # price / SMA
            mom_out = h_eom.iloc[-1,:].div(h_eom.iloc[[-t for t in range(1,self.LD+1)]].mean(axis=0),axis=0)
            mom_out = mom_out.loc[self.alldefensive].sort_values(ascending=False)
            # equal weightings to top relative momentum securities
            out_wts = pd.Series((1-pct_in)/self.TD,index=mom_out.index[:self.TD])
            wts = pd.concat([wts,out_wts])     
        
        self.Plot("PCT In", 'PCT', pct_in)

        wts = wts.groupby(wts.index).sum()

        return wts

    
    def OnEndOfDay(self):

        if not self.LiveMode:
            mkt_price = self.Securities[self.BNC].Close

        # the below fixes the divide by zero error in the MKT plot
        if mkt_price > 0 and mkt_price is not None:
            self.mkt.append(mkt_price)

        if len(self.mkt) >= 2 and not self.IsWarmingUp:
            mkt_perf = self.mkt[-1] / self.mkt[0] * self.cap
            self.Plot('Strategy Equity', self.BNC, mkt_perf)