#region imports
from AlgorithmImports import *
#endregion

class LiveReadyMomentum(QCAlgorithm):

    def Initialize(self):
        # 1. Timeline: Jan 2021 to Present
        self.SetStartDate(2020, 11, 1)  
        self.SetEndDate(2026, 1, 28)    
        self.SetCash(100000) 
        self.SetName("AI Momentum + 7% Trailing Stop")

        # 2. Universe: High Beta AI
        self.tickers = [
            "NVDA", "SMCI", "PLTR", "AMD", 
            "META", "AVGO", "TSM", "MSFT"
        ]
        
        self.symbols = []
        self.data = {}
        
        # Track the "High Water Mark" for each holding
        self.high_water_marks = {}

        for ticker in self.tickers:
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            self.symbols.append(symbol)
            # Faster Trend Filter (21 EMA) for quick trend detection
            self.data[symbol] = self.EMA(symbol, 21, Resolution.Daily)
            self.high_water_marks[symbol] = 0

        # 3. Parameters
        self.lookback = 63        # 3 Month Momentum
        self.hold_count = 2       # Top 2 Stocks
        self.stop_loss_pct = 0.07 # 7% Trailing Stop

        # 4. Schedule Rebalance (Weekly)
        self.Schedule.On(
            self.DateRules.WeekStart("NVDA"),
            self.TimeRules.AfterMarketOpen("NVDA", 30),
            self.Rebalance
        )
        
        self.SetBenchmark("SPY")
        self.SetSecurityInitializer(lambda x: x.SetFeeModel(ConstantFeeModel(1.0)))
        self.SetWarmUp(self.lookback + 10)

    def OnData(self, data):
        if self.IsWarmingUp: return

        # --- TRAILING STOP LOGIC ---
        # This runs every day to protect your gains
        for symbol in self.Portfolio.Keys:
            if not self.Portfolio[symbol].Invested:
                self.high_water_marks[symbol] = 0 # Reset if not owned
                continue

            # Check if we have data
            if not data.ContainsKey(symbol) or data[symbol] is None:
                continue

            current_price = data[symbol].Close

            # 1. Update High Water Mark
            # If price is higher than our record, update the record
            if current_price > self.high_water_marks[symbol]:
                self.high_water_marks[symbol] = current_price

            # 2. Check for Stop Loss Hit
            # If price is 7% below the High Water Mark -> SELL
            stop_price = self.high_water_marks[symbol] * (1 - self.stop_loss_pct)
            
            if current_price < stop_price:
                self.Liquidate(symbol, "Trailing Stop Hit")
                # Reset high water mark to prevent immediate rebuy issues
                self.high_water_marks[symbol] = 0 

    def Rebalance(self):
        if self.IsWarmingUp: return

        # 1. Calculate Momentum
        history = self.History(self.symbols, self.lookback, Resolution.Daily)
        if history.empty: return

        scores = {}
        for symbol in self.symbols:
            if symbol not in history.index: continue
            
            closes = history.loc[symbol]['close']
            if len(closes) < self.lookback: continue

            # Momentum Calculation
            start_price = closes.iloc[0]
            end_price = closes.iloc[-1]
            if start_price == 0: continue
            momentum = (end_price / start_price) - 1.0
            
            # TREND FILTER (Entry Requirement)
            # Only buy if price is > 21 EMA
            # AND if we haven't just been stopped out (Price > Stop Price check implied)
            ema = self.data[symbol].Current.Value
            if end_price > ema:
                scores[symbol] = momentum
            else:
                scores[symbol] = -100

        # 2. Select Winners
        sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        selected = [x[0] for x in sorted_scores[:self.hold_count] if x[1] > -100]

        # 3. Execution
        for symbol in self.Portfolio.Keys:
            # Don't sell if it's a winner we want to keep
            # UNLESS it hit the trailing stop (handled in OnData)
            if self.Portfolio[symbol].Invested and symbol not in selected:
                self.Liquidate(symbol)

        if len(selected) > 0:
            weight = 1.0 / len(selected)
            for symbol in selected:
                # Only buy if not already invested to avoid fee churn
                # (Or rebalance if you want exact equal weights)
                if not self.Portfolio[symbol].Invested:
                    self.SetHoldings(symbol, weight)
                    # Initialize High Water Mark for new trade
                    self.high_water_marks[symbol] = self.Securities[symbol].Price