# region imports
from AlgorithmImports import *
# endregion

class SimpleEqualWeightSP500Replicator(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100000)
        
        # Set rebalancing flag and date trackers
        self.rebalance_day = -1
        self.last_month = -1
        
        # Subscribe to SPY for benchmarking
        self.AddEquity("SPY", Resolution.Daily)
        self.SetBenchmark("SPY")
        
        # Set universe settings for better performance
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.ExtendedMarketHours = False
        
        # Add universe selection
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # Set trading warm-up period
        self.SetWarmUp(timedelta(days=30))

    def CoarseSelectionFunction(self, coarse):
        # Only run once per month
        if self.Time.month == self.last_month:
            return Universe.Unchanged
        
        self.last_month = self.Time.month
        self.rebalance_day = self.Time.day
        
        # Filter for stocks with fundamental data, price > $5, and good liquidity
        filtered = [x for x in coarse if x.HasFundamentalData 
                   and x.Price > 5 
                   and x.DollarVolume > 10000000]
        
        # Sort by dollar volume (liquid stocks)
        sorted_by_volume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        
        # Take top 1000 for fine selection
        return [x.Symbol for x in sorted_by_volume[:1000]]
    
    def FineSelectionFunction(self, fine):
        # Filter out stocks with negative earnings or extreme valuations
        filtered_fine = [x for x in fine if x.MarketCap > 3e9]
        
        # Sort by market cap in DESCENDING order (largest first)
        sorted_by_market_cap = sorted(filtered_fine, key=lambda x: x.MarketCap, reverse=True)
        
        # Take top 500 by market cap (or fewer if not enough stocks meet criteria)
        count = min(500, len(sorted_by_market_cap))
        selected = sorted_by_market_cap[:count]
        
        self.Log(f"Selected {len(selected)} stocks based on market cap")
        
        return [x.Symbol for x in selected]
    
    def OnData(self, data):
        # Skip if we're in the warm-up period
        if self.IsWarmingUp: 
            return
        
        # Only rebalance on the day we selected our universe
        if self.Time.day != self.rebalance_day: 
            return
        
        # Liquidate positions that are no longer in our universe
        for holding in self.Portfolio.Values:
            if holding.Invested and holding.Symbol not in self.ActiveSecurities:
                self.Liquidate(holding.Symbol)
        
        # Calculate equal weight for each position
        active_securities = [sec for sec in self.ActiveSecurities.Values 
                           if sec.Symbol.Value != "SPY" and sec.HasData]
        
        if len(active_securities) > 0:
            weight = 1.0 / len(active_securities)
            
            # Set holdings to equal weight only if security has data
            for security in active_securities:
                # Check if the security has valid price data
                if security.Price > 0 and security.IsTradable:
                    self.SetHoldings(security.Symbol, weight)
            
            self.Log(f"Rebalanced portfolio: {len(active_securities)} positions at {weight:.4f} weight each")