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.tickers = ["MSFT", "AAPL", "NVDA", "AMZN", "GOOG", "GOOGL", "META", "BRK/A", "BRK/B", "LLY", "AVGO", "TSLA", "JPM", "V", "WMT", "UNH", "MA", "XOM", "JNJ", "HD", "PG", "COST", "AMD", "MRK", "ABBV", "ORCL", "CRM", "CVX", "BAC", "NFLX", "KO", "ADBE", "PEP", "TMO", "LIN", "MCD", "ABT", "DIS"]
        self.symbols = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in self.tickers]
        self.AddEquity("SPY", Resolution.Daily)  # Benchmark
        
        for symbol in self.symbols:
            self.AddEquity(symbol.Value, Resolution.Daily)
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 10), self.RankStocks)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.BeforeMarketClose("SPY", 10), self.Rebalance)
        
        self.ranked_symbols = []

    def RankStocks(self):
        scores = {}
        for symbol in self.symbols:
            history = self.History(symbol, 210, Resolution.Daily)  # Fetching enough data for 200-day SMA
            
            if history.empty or len(history) < 210:
                self.Log(f"Not enough data for {symbol.Value} to rank")
                continue
            
            # Calculate Indicators
            close = history['close']
            momentum = (close.iloc[-1] / close.iloc[-90]) - 1
            rsi = self.RSI(symbol, 14, MovingAverageType.Wilders, Resolution.Daily, Field.Close).Current.Value
            sma50 = self.SMA(symbol, 50, Resolution.Daily, Field.Close).Current.Value
            sma200 = self.SMA(symbol, 200, Resolution.Daily, Field.Close).Current.Value
            
            # Example scoring formula
            score = momentum + rsi/100 + (sma50 + sma200)/2
            scores[symbol] = score
        
        # Sort symbols by score in descending order and keep the top 10
        self.ranked_symbols = sorted(scores, key=scores.get, reverse=True)[:10]

    def Rebalance(self):
        if len(self.ranked_symbols) == 0:
            self.Log("No ranked symbols to rebalance")
            return
        
        # Calculate portfolio target size for each position
        target_size = 1.0 / len(self.ranked_symbols)
        
        # Sell symbols not in the top 10 anymore
        current_holdings = [s.Symbol for s in self.Portfolio.Values if s.Invested and s.Symbol not in self.ranked_symbols]
        for symbol in current_holdings:
            self.Liquidate(symbol)
        
        # Buy or adjust positions for the top 10 stocks
        for symbol in self.ranked_symbols:
            self.SetHoldings(symbol, target_size)
    
    def OnData(self, data):
        # If you don't need to do anything on each data slice, you can leave this empty
        pass