# region imports
from AlgorithmImports import *
import datetime
# endregion

class NdxMonthly(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 12, 1)
        # self.SetEndDate(2023, 12, 31)
        self.SetEndDate(datetime.date.today())
        self.SetCash(100000)
        
        self.SetSecurityInitializer(lambda security: security.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted))

        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.spx = self.AddIndex("SPX", Resolution.Daily).Symbol
        self.spy_ema5 = self.EMA(self.spy, 5)   
        self.spy_sma200 = self.SMA(self.spy, 200)
        
        # set data normaliaztion mode
        
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.universe_etf = self.AddUniverse(self.Universe.ETF("QQQ", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))

        self.EnableAutomaticIndicatorWarmUp = True

        self.indicators = {}

    def ETFConstituentsFilter(self, constituents: List[ETFConstituentData]) -> List[Symbol]:
        indicators = {}

        for c in constituents:  
            symbol = c.Symbol
  
            if c not in self.ActiveSecurities.Keys:
                self.AddEquity(symbol, Resolution.Daily)          
                            
            roc_short = self.ROC(symbol, 120).Current.Value * 100
            roc_long = self.ROC(symbol, 240).Current.Value * 100
            sma120 = self.SMA(symbol, 120)
            ema10 = self.EMA(symbol, 10)
            momentum = roc_short * (252/120) * .95 + roc_long * (252/240) * 1.05  # slightly weighted towards long

            factor = 0.4 * self.ROC(symbol, 63).Current.Value * 100 + 0.2 * self.ROC(symbol, 126).Current.Value * 100 + 0.2 * self.ROC(symbol, 189).Current.Value * 100 + 0.2 * self.ROC(symbol, 252).Current.Value * 100
            
            # Add the indicators to the dictionary
            indicators[symbol] = round(momentum, 2)
            
        # remove securities that we are subsribed to but are not in constituents
        for symbol in self.ActiveSecurities.Keys:
            if symbol not in indicators:
                self.RemoveSecurity(symbol)
            
        self.indicators = {c.Symbol: indicators[c.Symbol] for c in constituents}

        return [c.Symbol for c in constituents]

    def OnData(self, data: Slice) -> None:
        if self.IsWarmingUp:
            # log we are in warmup period
            self.Log("Warming up...")
            return

        bull_mkt = self.spy_ema5.Current.Value > self.spy_sma200.Current.Value
        if bull_mkt: # and self.Time.day == 1:
            self.Log("Bull market - number of universe members are currently " + str(len(self.indicators)))   
        else:
            self.Log(f'Not bull market.  SPY EMA5: {self.spy_ema5.Current.Value}, SPY SMA200: {self.spy_sma200.Current.Value}')
            # for symbol, indicator in self.indicators.items():
            #     self.Log(f"{symbol}: monthly momo {indicator['momentum']}")
            
            # sort the universe by momentum
        sorted_dict = dict(sorted(self.indicators.items(), key=lambda x: x[1], reverse=True))
        symbol_dict = {symbol.Value: value for symbol, value in sorted_dict.items()}
        # Only keep the top 20 items in symbol_dict
        symbol_dict = dict(list(symbol_dict.items())[:20])
        self.Log("Sorted universe by momentum")
        self.Log(str(symbol_dict))

        str_output = ''
        for symbol, indicator in symbol_dict.items():
            str_output += f"{symbol},{indicator}\n"

        try:
            spx_history = self.History([self.spx], 1, Resolution.Daily)
            spx_close = spx_history.loc[self.spx].close[-1]
            self.Log(f"SPX close price: {spx_close}")
            self.Log(str(spx_history))
        except:
            self.Log("Error getting SPX close price")