# https://quantpedia.com/strategies/asset-class-momentum-rotational-system/
#
# Use 5 ETFs (SPY - US stocks, EFA - foreign stocks, IEF - 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.

#region imports
from AlgorithmImports import *
#endregion

class MomentumAssetAllocationStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetCash(100000)
        
        self.data:dict[str, RateOfChange] = {}
        period:int = 12 * 21
        self.SetWarmUp(period, Resolution.Daily)
        
        self.traded_count:int = 3

        self.symbols:List[str] = ["SPY", "QQQ", "FXI", "USO", "GLD"]

        for symbol in self.symbols:
            self.AddEquity(symbol, Resolution.Minute)
            self.data[symbol] = self.ROC(symbol, period, Resolution.Daily)
        
        self.recent_month:int = -1

    def OnData(self, data):
        if self.IsWarmingUp: return
        
        if not (self.Time.hour == 9 and self.Time.minute == 31):
            return

        self.Log(f"Market Open Time: {self.Time}")

        # rebalance once a month
        if self.Time.month == self.recent_month:
            return
        self.recent_month = self.Time.month
        
        self.Log(f"New monthly rebalance...")

        # debug/log info
        selected:dict[str, RateOfChange] = {}
        
        for symbol, roc in self.data.items():
            data_ready:bool = bool(symbol in data and data[symbol])
            roc_ready:bool = bool(roc.IsReady)

            self.Log(f"Data for {symbol} are present: {data_ready}")
            self.Log(f"ROC for {symbol} IsReady: {roc_ready}")

            if data_ready and roc_ready:
                selected[symbol] = roc

        sorted_by_momentum:List = sorted(selected.items(), key = lambda x: x[1].Current.Value, reverse = True)
        # sorted_by_momentum = sorted([x for x in self.data.items() if x[1].IsReady and x[0] in data and data[x[0]]], key = lambda x: x[1].Current.Value, reverse = True)
        self.Log(f"Number of assets to sort: {len(sorted_by_momentum)}; at least {self.traded_count} needed.")

        long:List[str] = []
        
        if len(sorted_by_momentum) >= self.traded_count:
            long = [x[0] for x in sorted_by_momentum][:self.traded_count]

        invested:List[str] = [x.Key.Value for x in self.Portfolio if x.Value.Invested]
        for symbol in invested:
            if symbol not in long:
                self.Liquidate(symbol)

        self.Log(f"Selected long leg for next month: {long}")
        for symbol in long:
            self.SetHoldings(symbol, 1 / len(long))