#region imports
from AlgorithmImports import *
#endregion

class BubbleSurfer(QCAlgorithm):

    def Initialize(self):
        # 1. Timeline: Start of the AI Mania (2020)
        self.SetStartDate(2020, 1, 1)  
        self.SetEndDate(2026, 1, 28)    
        self.SetCash(100000) 
        self.SetName("AI Bubble Surfer (Chandelier Exit)")

        # 2. Universe: The "AI Bubble" Basket
        self.tickers = [
            "NVDA", "SMCI", "AMD", "TSM", 
            "AVGO", "META", "MSFT", "PLTR", "GOOGL"
        ]
        
        self.data = {}
        
        for ticker in self.tickers:
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            
            # 3. Indicators for the Chandelier Exit
            # ATR (Average True Range) measures the "noise" or "heat" of the stock
            atr = self.ATR(symbol, 14, MovingAverageType.Wilders, Resolution.Daily)
            
            # MAX indicator to track the "Highest High" over the lookback period
            highest_high = self.MAX(symbol, 22, Resolution.Daily, Field.High)
            
            self.data[symbol] = {
                'atr': atr,
                'max': highest_high,
                'stop_price': 0
            }

        # 4. Parameters
        self.atr_mult = 3.0       # 3x ATR is the standard "Trend Following" distance
        self.lookback = 63        # 3 Month Momentum ranking
        
        self.current_winner = None # Track our single holding

        # Check regularly
        self.Schedule.On(
            self.DateRules.EveryDay("NVDA"),
            self.TimeRules.AfterMarketOpen("NVDA", 30),
            self.ManagePositions
        )
        
        self.SetBenchmark("SPY")
        self.SetSecurityInitializer(lambda x: x.SetFeeModel(ConstantFeeModel(1.0)))
        self.SetWarmUp(100)

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

        # 1. Update Stop Prices (The Chandelier Logic)
        for symbol, indicators in self.data.items():
            if not indicators['atr'].IsReady or not indicators['max'].IsReady:
                continue
            
            # Chandelier Exit Formula:
            # Stop = Highest High of last 22 days - (3 * Current ATR)
            # This "hangs" the stop loss from the highest point reached.
            hh = indicators['max'].Current.Value
            volatility = indicators['atr'].Current.Value
            
            new_stop = hh - (volatility * self.atr_mult)
            
            # Trailing Logic: Stop can ONLY move UP, never down (ratchet)
            # Unless we are not invested, then it resets.
            if self.Portfolio[symbol].Invested:
                if new_stop > indicators['stop_price']:
                    indicators['stop_price'] = new_stop
            else:
                indicators['stop_price'] = new_stop

        # 2. CHECK EXITS (Risk Management First)
        # If our current winner hit the floor, SELL immediately.
        if self.current_winner and self.Portfolio[self.current_winner].Invested:
            current_price = self.Securities[self.current_winner].Price
            stop_price = self.data[self.current_winner]['stop_price']
            
            if current_price < stop_price:
                self.Liquidate(self.current_winner, "Chandelier Exit Hit")
                self.current_winner = None # We are now Cash

        # 3. CHECK ENTRIES (FOMO Logic)
        # Only look for a new trade if we are 100% Cash (Winner takes all)
        if not self.Portfolio.Invested:
            self.EnterTopStock()

    def EnterTopStock(self):
        # Rank universe by 3-Month Momentum
        history = self.History(list(self.data.keys()), self.lookback, Resolution.Daily)
        if history.empty: return

        scores = {}
        for symbol in self.data.keys():
            if symbol not in history.index: continue
            closes = history.loc[symbol]['close']
            if len(closes) < self.lookback: continue
            
            # Momentum
            mom = (closes.iloc[-1] / closes.iloc[0]) - 1
            scores[symbol] = mom

        # Sort and Pick #1
        if not scores: return
        best_symbol = sorted(scores.items(), key=lambda x: x[1], reverse=True)[0][0]
        
        # FINAL SAFETY CHECK:
        # Don't buy if it's currently crashing (Price < Calculated Stop)
        # This prevents buying a "falling knife" even if it has high past momentum.
        price = self.Securities[best_symbol].Price
        stop = self.data[best_symbol]['stop_price']
        
        if price > stop:
            self.SetHoldings(best_symbol, 1.0)
            self.current_winner = best_symbol
            # Initialize the stop strictly
            self.data[best_symbol]['stop_price'] = self.data[best_symbol]['max'].Current.Value - (self.data[best_symbol]['atr'].Current.Value * self.atr_mult)