# The investment universe consists of the Offensive: SPY, QQQ, IWM, VGK, EWJ, VWO, VNQ, DBC, GLD, TLT, HYG, LQD; Protective: SPY, VWO, VEA, BND; and the Defensive: TIP, DBC, BIL, IEF,
# TLT, LQD, BND. The trading algorithm is: on the close of the last trading day of each month t. 
# 1. Calculate a relative momentum score for each of assets in the offensive and defensive universe, where relative momentum at t equals pt / SMA(12) – 1. Note that the slow SMA(12)
# trend is calculated based on month-end values with maximum lag 12, so as the average over pt. pt-12 representing the most recent 13 month-end prices, including today.
# 2. Select the Top3 from a defensive universe with both relative and absolute SMA(12) momentum, if at least one of the assets in the protective (or canary) universe show negative 
# absolute momentum, where absolute momentum at t is based on fast momentum 13612W, which is the weighted average of returns over 1, 3, 6 and 12 months with weights 12, 4, 2, 1, resp.
# Otherwise, select the offensive universe. 
# 3. Depending on step 2, select the Top 6 assets with the highest relative momentum value of the offensive or the defensive universe and allocate 1/(Top 6) of the portfolio to each.
# Replace the ‘bad’ defensive selections (assets with momentum less than BIL) by BIL. Hold positions until the final trading day of the following month. Re-balance the entire portfolio
# monthly, regardless of whether there is a change in position. This implies that the switching is 100% to defensive when at least one of the Protective (or ‘canary’) universe assets
# shows negative (or ‘bad’) absolute momentum, no switching (so 0% defensive) with no canary assets ‘bad’.

# region imports
from AlgorithmImports import *
import pandas as pd
import numpy as np
from typing import List, Dict
from pandas.core.frame import DataFrame
# endregion

class BoldAssetAllocation(QCAlgorithm):
    
    def Initialize(self):
        self.SetCash(100000)
        # self.SetStartDate(2008, 1, 1)
        # self.SetStartDate(2023, 4, 1)  # Set Start Date
        # self.SetEndDate(2023, 6, 1)  # Set End Date
        self.SetStartDate(2018, 1, 1)  # Set Start Date
        self.SetEndDate(2024, 8, 31)  # Set End Date

        # all assets
        self.offensive:List[str] = [
            "SPY", "QQQ",
            "IWM", "VGK",
            "EWJ", "VWO",
            "VNQ", "DBC",
            "GLD", "TLT",
            "HYG", "LQD",
        ]

        # Max ETF Fees if chose the 6 most expensive comes to 0.44%
        # SPY: .09%
        # QQQ: .20% 5
        # IWM: .19% 6
        # VGK: .11%
        # EWJ: .50% 2
        # VWO: .08%
        # VNQ: .12%
        # DBC: .85% 1
        # GLD: .40% 4
        # TLT: .15%
        # HYG: .48% 3
        # LQD: .14%


        self.protective:List[str] = ["SPY", "VWO", "VEA", "BND"]
        self.defensive:List[str] = ["TIP", "DBC", "BIL", "IEF", "TLT", "LQD", "BND"]
        self.safe:str = "BIL"

        # Max ETF Fees if chose the 3 most expensive comes to 0.20%
        # TIP: .19% 2
        # DBC: .85% 1
        # BIL: .1354%
        # IEF: .15% 3
        # TLT: .15%
        # LQD: .14%
        # BND: .03%

        # strategy parameters (our implementation)
        self.prds:List[int] = [1, 3, 6, 12]                     # fast momentum settings
        self.prdweights:np.ndarray = np.array([12, 4, 2, 1])    # momentum weights
        self.LO, self.LP, self.LD, self.B, self.TO, self.TD = [
            len(self.offensive),
            len(self.protective),
            len(self.defensive),
            1,
            6,
            3,
        ]  # number of offensive, protective, defensive assets, threshold for "bad" assets, select top n of offensive and defensive assets
        
        self.hprd:int = (max(self.prds + [self.LO, self.LD]) * 21 + 50)  # momentum periods calculation

        # repeat safe asset so it can be selected multiple times
        self.all_defensive:List[str] = self.defensive + [self.safe] * max(
            0, self.TD - sum([1 * (e == self.safe) for e in self.defensive])
        )
        
        self.equities:List[str] = list(
            dict.fromkeys(self.protective + self.offensive + self.all_defensive)
        )
        
        leverage:int = 3
        for equity in self.equities:
            data:Equity = self.AddEquity(equity, Resolution.Daily)
            data.SetLeverage(leverage)

        self.recent_month:int = -1

        self.settings.daily_precise_end_time = False

    def OnData(self, data:Slice) -> None:
        if self.IsWarmingUp:
            return

        # monthly rebalance
        if self.recent_month == self.Time.month:
            return
        self.recent_month = self.Time.month

        # get price data and trading weights
        h:DataFrame = self.History(self.equities, self.hprd, Resolution.Daily)["close"].unstack(level=0)
        weights:pd.Series = self.trade_weights(h)

        # # Used to display the start and end date of historical price dataframe
        # self.Debug("Date 1: " + str(h.iloc[:1].index))
        # self.Debug("Date 2: " + str(h.iloc[-1:].index))

        # trade
        self.SetHoldings([PortfolioTarget(x, y) for x, y in zip(weights.index, weights.values) if x in data and data[x]])

    def trade_weights(self, hist:DataFrame) -> pd.Series:
        # initialize weights series
        weights:pd.Series = pd.Series(0, index=hist.columns)
        # end of month values
        h_eom:DataFrame = hist.loc[hist.groupby(hist.index.to_period("M")).apply(lambda x: x.index.max())].iloc[:-1, :]

        # Check if protective universe is triggered.

        # build dataframe of momentum values
        mom:DataFrame = (h_eom.iloc[-1, :].div(h_eom.iloc[[-p - 1 for p in self.prds], :], axis=0) - 1)
        mom = mom.loc[:, self.protective].T

        # determine number of protective securities with negative weighted momentum
        n_protective:float = np.sum(np.sum(mom.values * self.prdweights, axis=1) < 0)
        
        # % equity offensive
        pct_in:float = 1 - min(1, n_protective / 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_weights = pd.Series(pct_in / self.TO, index=mom_in.index[:self.TO])
            weights = pd.concat([weights, in_weights])
        
        # 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.all_defensive].sort_values(ascending=False)
            # equal weightings to top relative momentum securities
            out_weights = pd.Series((1 - pct_in) / self.TD, index=mom_out.index[:self.TD])
            weights = pd.concat([weights, out_weights])

        weights:pd.Series = weights.groupby(weights.index).sum()

        return weights