from AlgorithmImports import *

class SMAStrategy(QCAlgorithm):

    def Initialize(self):
        # --- Strategy Initialization ---
        self.SetStartDate(2010, 9, 1)
        self.SetEndDate(2024, 9, 1)
        self.SetCash(100000)  # Set the initial capital for the strategy

        # --- Symbol Selection ---
        # We will use QQQ for the SMA calculations, but trade TQQQ and SQQQ
        self.qqq = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.tqqq = self.AddEquity("TQQQ", Resolution.Daily).Symbol
        self.sqqq = self.AddEquity("SQQQ", Resolution.Daily).Symbol

        # --- Data Preparation: SMA Definitions ---
        # Define fast and slow SMAs for QQQ
        self.fast_sma = self.SMA(self.qqq, 100, Resolution.Daily)  # Fast SMA (around 5 months)
        self.slow_sma = self.SMA(self.qqq, 200, Resolution.Daily)  # Slow SMA (around 10 months)

        # Initialize benchmark for QQQ Buy and Hold
        self.benchmark_initial_price = None
        self.benchmark_initial_cash = self.Portfolio.Cash  # Start with the same cash amount for benchmark

        # --- Data Warmup: Ensure enough data for SMA calculations ---
        self.SetWarmUp(252)  # Warm up period of 252 days (~12 months)

    def OnData(self, data: Slice):
        # --- Data Preparation and Validity Check ---
        # Ensure data is warmed up and available
        if self.IsWarmingUp:
            return
        
        # Ensure data is available for QQQ (used for SMA and benchmark)
        if not data.ContainsKey(self.qqq) or not data[self.qqq]:
            return

        # --- Benchmark Plotting: QQQ Buy and Hold Calculation ---
        if self.benchmark_initial_price is None:
            # Initialize the benchmark with the first price available
            self.benchmark_initial_price = data[self.qqq].Price
        
        # Calculate the current value of the QQQ Buy and Hold benchmark
        benchmark_return = (data[self.qqq].Price - self.benchmark_initial_price) / self.benchmark_initial_price
        benchmark_value = self.benchmark_initial_cash * (1 + benchmark_return)
        
        # Plot the QQQ Buy and Hold benchmark
        self.Plot("Strategy Equity", "QQQ Buy and Hold", benchmark_value)

        # --- Trade Logic: Buy TQQQ or SQQQ based on SMA ---
        # Ensure SMAs are ready before executing trades
        if not self.fast_sma.IsReady or not self.slow_sma.IsReady:
            return

        # If fast SMA > slow SMA, buy TQQQ (leveraged long)
        if self.fast_sma.Current.Value > self.slow_sma.Current.Value:
            # Sell SQQQ if held
            if self.Portfolio[self.sqqq].Invested:
                self.Liquidate(self.sqqq)
            # Buy TQQQ if not already invested
            if not self.Portfolio[self.tqqq].Invested:
                self.SetHoldings(self.tqqq, 1.0)  # Invest 100% in TQQQ
        
        # If fast SMA <= slow SMA, sell QQQ and buy SQQQ (leveraged short)
        elif self.fast_sma.Current.Value <= self.slow_sma.Current.Value:
            # Liquidate TQQQ if held
            if self.Portfolio[self.tqqq].Invested:
                self.Liquidate(self.tqqq)
            # Buy SQQQ if not already invested
            if not self.Portfolio[self.sqqq].Invested:
                self.SetHoldings(self.sqqq, 1.0)  # Invest 100% in SQQQ