#region imports
from AlgorithmImports import *
#endregion
class VerticalNadionShield(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)  # Set Start Date
        self.SetEndDate(2022, 12, 31)  # set end date
        self.SetCash(100000)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.leverage = 1
        
  
        self.equities = ["SPY", "SCZ"]
        self.equityCombinedMomentum = {}
        
        self.bonds = ["TLT", "TIP"]
        self.bondCombinedMomentum = {}

            
        for equity in self.equities:
            self.AddEquity(equity, Resolution.Hour)
            self.Securities[equity].SetDataNormalizationMode(DataNormalizationMode.TotalReturn)
            self.equityCombinedMomentum[equity] = CombinedMomentum(self, equity)

        for bond in self.bonds:
            self.AddEquity(bond, Resolution.Hour)
            self.Securities[bond].SetDataNormalizationMode(DataNormalizationMode.TotalReturn)
            self.bondCombinedMomentum[bond] = CombinedMomentum(self, bond)
            
        self.SetWarmUp(125)

    def shiftAssets(self, target):
        if not (self.Portfolio[target].Invested):
            for symbol in self.Portfolio.Keys:
                self.Liquidate(symbol)
            if not self.Portfolio.Invested:
                self.MarketOnCloseOrder(target, self.CalculateOrderQuantity(target, 1 * self.leverage))


    def getMonthLastTradingDay(self):
        month_last_day = DateTime(self.Time.year, self.Time.month, DateTime.DaysInMonth(self.Time.year, self.Time.month))
        tradingDays = self.TradingCalendar.GetDaysByType(TradingDayType.BusinessDay, DateTime(self.Time.year, self.Time.month, 1), month_last_day)
        tradingDays = [day.Date.date() for day in tradingDays]
        return tradingDays[-1]
        

    def OnData(self, data):
        if self.IsWarmingUp:
            return

        if (self.Time.date() == self.getMonthLastTradingDay()) and (self.Time.hour == 15):
            topEquities = sorted(self.equityCombinedMomentum.items(), key=lambda x: x[1].getValue(), reverse=True)
            topBonds = sorted(self.bondCombinedMomentum.items(), key=lambda x: x[1].getValue(), reverse=True)
            if (topEquities[0][1].getValue() > 0):
                self.shiftAssets(topEquities[0][0])
            else:
                self.shiftAssets(topBonds[0][0])
        

class CombinedMomentum():
    def __init__(self, algo, symbol):
        self.fst = algo.MOMP(symbol,  21, Resolution.Daily)
        self.med = algo.MOMP(symbol,  63, Resolution.Daily)
        self.slw = algo.MOMP(symbol,  126, Resolution.Daily)
        
    def getValue(self):
        value = (self.fst.Current.Value + self.med.Current.Value + self.slw.Current.Value) / 3
        return value