from AlgorithmImports import *

class MomentumAndSMAStrategy(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2014, 1, 1)  # Start Date
        self.SetEndDate(2024, 1, 1)  # End Date
        self.SetCash(10000)  # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.numberOfSymbols = 10
        self.ranked_symbols = []

        # Add SPY for scheduling purpose, ensuring it's available in the algorithm
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.Schedule.On(self.DateRules.EveryDay(self.spy), self.TimeRules.BeforeMarketClose(self.spy, 10), self.RankAndRebalance)

    def CoarseSelectionFunction(self, coarse):
        filtered = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)[:100]
        return [c.Symbol for c in filtered if c.HasFundamentalData]

    def FineSelectionFunction(self, fine):
        return [f.Symbol for f in fine if f.MarketCap > 2e9][:self.numberOfSymbols]

    def RankAndRebalance(self):
        if self.Time.weekday() != 0:  # Only run on Mondays
            return
        
        self.liquidStocks = self.ranked_symbols
        
        scores = {}
        for symbol in self.liquidStocks:
            history = self.History(symbol, 210, Resolution.Daily)
            if history.empty or len(history) < 90:
                self.Log(f"Not enough data for {symbol.Value}")
                continue
            
            close = history['close']
            momentum = (close.iloc[-1] / close.iloc[-90]) - 1
            rsi = self.RSI(symbol, 14, Resolution.Daily).Current.Value
            sma50 = self.SMA(symbol, 50, Resolution.Daily).Current.Value
            sma200 = self.SMA(symbol, 200, Resolution.Daily).Current.Value
            
            score = momentum + rsi/100 + (sma50 + sma200)/2
            scores[symbol] = score
        
        self.ranked_symbols = sorted(scores, key=scores.get, reverse=True)[:self.numberOfSymbols]
        
        if not self.ranked_symbols:
            self.Log("No symbols to rank/rebalance.")
            return
        
        for holding in self.Portfolio.Values:
            if holding.Invested and holding.Symbol not in self.ranked_symbols:
                self.Liquidate(holding.Symbol)
        
        for symbol in self.ranked_symbols:
            self.SetHoldings(symbol, 1 / len(self.ranked_symbols))

    def OnSecuritiesChanged(self, changes):
        self.ranked_symbols = [c.Symbol for c in changes.AddedSecurities]