#region imports
from AlgorithmImports import *

import pandas as pd
from datetime import datetime
#endregion
# https://quantpedia.com/Screener/Details/2
# Use 5 ETFs (SPY - US stocks, EFA - foreign stocks, BND - bonds, VNQ - REITs, GSG - commodities).
# Pick 3 ETFs with strongest 12 month momentum into your portfolio and weight them equally. 
# Hold for 1 month and then rebalance.


class AssetClassMomentumAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2007, 5, 1) 
        self.set_cash(100000) 
        # create a dictionary to store momentum percent indicators for all symbols 
        self._data = {}
        period = 12*21
        symbols = ["SPY", "EFA", "BND", "VNQ", "GSG"]
        # warm up the MOMP indicator
        self.set_warm_up(period+1)
        for symbol in symbols:
            self.add_equity(symbol, Resolution.DAILY)
            self._data[symbol] = self.momp(symbol, period, Resolution.DAILY)
        # shcedule the function to fire at the month start 
        self.schedule.on(self.date_rules.month_start("SPY"), self.time_rules.after_market_open("SPY"), self._rebalance)

    def _rebalance(self):
        if self.is_warming_up: 
            return
        top3 = pd.Series(self._data).sort_values(ascending=False)[:3]
        for symbol, security_hold in self.portfolio.items():
            # liquidate the security which is no longer in the top3 momentum percent list
            if security_hold.invested and (symbol.value not in top3.index):
                self.liquidate(symbol)
        
        added_symbols = []        
        for symbol in top3.index:
            if not self.portfolio[symbol].invested:
                added_symbols.append(symbol)
        for added in added_symbols:
            self.set_holdings(added, 1/len(added_symbols))