#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 = ["XLV", "XLK", "XLI", "XLU", "XLF", "XLY", "XLP", "XLB", "XLE", "PSR", "IYZ", "USO", "SCZ", "SH", "PSQ", "QQQ"]
        self.equityCombinedMomentum = {}
        
        self.bonds = ["TLT", "TIP", "BIL", "AGG", "HYG"]
        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):
        for symbol in self.Portfolio.Keys:
                if symbol.Value not in [x[0] for x in target]:
                    self.Liquidate(symbol)
        for x in [x[0] for x in target]:
            if not (self.Portfolio[x].Invested):
                self.MarketOnCloseOrder(x, self.CalculateOrderQuantity(x, 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

        print(self.equities)

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

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