#region imports
from AlgorithmImports import *
#endregion

class AIBasketStrategy(QCAlgorithm):

    def Initialize(self):
        # 1. Setup: Last 5 Years
        self.SetStartDate(2020, 11, 10)  
        self.SetEndDate(2026, 1, 1)    
        self.SetCash(100000) 

        # 2. Define the "AI Basket"
        # We manually select the leaders to ensure purity of the theme
        self.tickers = [
            "NVDA", # GPU King
            "AMD",  # GPU Challenger
            "TSM",  # The Foundry (makes the chips)
            "AVGO", # Broadcom (Custom AI chips & Networking)
            "MSFT", # OpenAI / Cloud
            "GOOGL",# Gemini / DeepMind
            "META", # LLaMA / Meta AI
            "PLTR"  # AI Data Foundry
        ]
        
        self.symbols = []
        for ticker in self.tickers:
            # Add equity and set resolution to Daily (sufficient for rebalancing)
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            self.symbols.append(symbol)

        # 3. Schedule Rebalancing (First trading day of every month)
        # This keeps the portfolio strictly equal-weighted
        self.Schedule.On(
            self.DateRules.MonthStart("NVDA"),
            self.TimeRules.AfterMarketOpen("NVDA", 30),
            self.Rebalance
        )
        
        # 4. Realistic Settings
        self.SetSecurityInitializer(lambda x: x.SetFeeModel(ConstantFeeModel(1.0)))
        self.SetBenchmark("SPY") # Compare against S&P 500

    def Rebalance(self):
        # Calculate target weight (100% / 8 stocks = 12.5% each)
        # We check if stock has data to avoid errors (e.g. if one was delisted or halted)
        valid_symbols = [s for s in self.symbols if self.Securities[s].Price > 0]
        
        if len(valid_symbols) == 0: return
        
        weight = 1.0 / len(valid_symbols)
        
        # Execute Trades
        for symbol in valid_symbols:
            self.SetHoldings(symbol, weight)
            
    def OnData(self, data):
        # We don't need trade logic here because the Scheduler handles it.
        pass