#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):
        self.SetStartDate(2008, 1, 1)
        self.start_cash = 100000
        self.SetCash(self.start_cash)
        
        self.SetBenchmark('SPY')
        
        self.leverage = 3

        # 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,1,3]
        self.hprd = max(self.prds+[self.LO,self.LD])*21+50
        
        # Assets
        self.canary = ['SPY','EFA','EEM','BND']
        self.offensive = ['QQQ','EFA','EEM','BND']
        self.defensive = ['BIL','BND','DBC','IEF','LQD','TIP','TLT']
        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:
            data = self.AddEquity(eq, Resolution.Minute)
            data.SetLeverage(self.leverage * 2)
        
        # monthly rebalance
        self.Schedule.On(self.DateRules.MonthStart(self.canary[0]),self.TimeRules.AfterMarketOpen(self.canary[0],30),self.rebal)
        self.Trade = True

        # benchmark stuff
        # self.benchmark_symbol:Symbol = self.AddEquity('TQQQ', Resolution.Daily).Symbol
        # self.benchmark_values = []
        # self.Schedule.On(self.DateRules.EveryDay(self.benchmark_symbol), self.TimeRules.BeforeMarketClose(self.benchmark_symbol, 0), self.update_eq_chart)

    def update_eq_chart(self):
        ''' Updates benchmark eqity in main Equity chart '''
        
        hist:df = self.History([self.benchmark_symbol], 2, Resolution.Daily)
        if not hist.empty:
            hist = hist['close'].unstack(level= 0).dropna() 
            self.benchmark_values.append(hist[self.benchmark_symbol].iloc[-1])
            benchmark_perf = self.benchmark_values[-1] / self.benchmark_values[0] * self.start_cash

            self.Plot("Strategy Equity", self.benchmark_symbol.Value, benchmark_perf)

    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*self.leverage) 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])     
        
        wts = wts.groupby(wts.index).sum()

        return wts