# Pick 5 ETFs with strongest 3 month momentum into your portfolio and weight them equally. 
# Hold for 1 month and then rebalance at the satrt of a month

import pandas as pd
from datetime import datetime

class AssetClassMomentumAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2002, 7, 28)  
        self.SetEndDate(datetime.now())  
        self.SetCash(100000) 
        # create a dictionary to store momentum indicators for all symbols 
        self.data = {}
        period = 3*22*24 # period of the momentum indicator is 3 months and average 22 business day for each month
        self.symbols = ["SPY", "IWF", "QQQ", "IBB", "XLV", "EEM", "IWD", "TLT", "IEF", "IGOV", "GLD", "IYR", "DBC"]
        
        
        # warm up the MOM indicator
        self.SetWarmUp(period)
        for symbol in self.symbols:
            self.AddEquity(symbol, Resolution.Hour)
            self.data[symbol] = self.MOM(symbol, period, Resolution.Hour)
            
        # schedule the function to fire at the month start, 30 min from a market open
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 30), self.Rebalance)
            
    def OnData(self, data):
        pass

    def Rebalance(self):
        if self.IsWarmingUp: return
        
        
        top5 = pd.Series(self.data).sort_values(ascending = False)[:5]
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            # liquidate the security which is no longer in the top5 momentum list
            if security_hold.Invested and (security_hold.Symbol.Value not in top5.index):
                self.Liquidate(security_hold.Symbol)
        
        added_symbols = []        
        for symbol in top5.index:
            if not self.Portfolio[symbol].Invested:
                added_symbols.append(symbol)
        
        # 0.9 is to have some cushion and enough capital for order execution    
        for added in added_symbols:
            self.SetHoldings(added, 0.9/len(added_symbols))