'''
Hybrid Asset Allocation by Wouter Keller
See:  https://papers.ssrn.com/sol3/papers.cfm?abstract_id=4346906
This is the "balanced" version from the paper
'''

# region imports
from AlgorithmImports import *
import numpy as np
# endregion

class HAA(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)  # Set Start Date
        # self.SetEndDate(2022, 12, 31)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        # Asset lists
        self.selO = ['SPY','IWM','VWO','VEA','VNQ','DBC','IEF','TLT']
        self.selD = ['BIL', 'IEF']
        self.selP = ['TIP']
        self.safe = ['BIL']
        # Add assets
        self.eqs = list(set(self.selO+self.selD+self.selP+self.safe))
        for eq in self.eqs:
            self.AddEquity(eq,Resolution.Minute)

        # Algo Parameters
        self.TO, self.TD = [4,1]
        self.mom_prds = [1,3,6,12]  #1, 3, 6, and 12 months.
        self.hprd = np.max(self.mom_prds)*35

        # Scheduler
        self.Schedule.On(self.DateRules.MonthStart(self.selO[0]),
                        self.TimeRules.AfterMarketOpen(self.selO[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.mom_prds],:],axis=0)-1
        mom_avg = mom.mean(axis=0)
        # Determine number of canary securities with negative momentum
        n_canary = np.sum((mom_avg[self.selP]<0)*1)/len(self.selP)
        # % equity offensive 
        pct_in = 1-n_canary

        # cash return
        pct_safe = 0
        # cash_r = mom_avg[self.safe][0]

        # =====================================
        # get weights for offensive and defensive universes
        # =====================================
        # determine weights of offensive universe
        if pct_in > 0:
            # Avg MOM
            mom_in = mom_avg.loc[self.selO].sort_values(ascending=False).iloc[:self.TO]
            mom_in = mom_in[mom_in>0]
            # equal weightings to top relative momentum securities
            in_wts = pd.Series(pct_in/self.TO,index=mom_in.index)
            pct_safe = 1-mom_in.shape[0]/self.TO
            wts = pd.concat([wts,in_wts])
        # determine weights of defensive universe
        if pct_in < 1:
            # Avg MOM
            mom_out = mom_avg.loc[self.selD].sort_values(ascending=False).iloc[:self.TD]
            # equal weightings to top relative momentum securities
            out_wts = pd.Series((1-pct_in)/self.TD,index=mom_out.index)
            wts = pd.concat([wts,out_wts])     
        
        # cash portion
        wts[wts.index.str.startswith(self.safe[0])] += pct_safe

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

        return wts