from AlgorithmImports import *
import numpy as np

class VarSwitchStrategyA(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetCash(10000)

        # Add assets
        self.qqq = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.tqqq = self.AddEquity("TQQQ", Resolution.Daily).Symbol

        # Set resolution and warm-up
        self.lookback = 252
        self.z = 2.58
        self.holding_period = 10
        self.window = RollingWindow[float](self.lookback)
        self.SetWarmUp(self.lookback + self.holding_period, Resolution.Daily)

        # Schedule check
        self.Schedule.On(
            self.DateRules.EveryDay(self.qqq),
            self.TimeRules.AfterMarketOpen(self.qqq, 1),
            self.DailyCheck
        )

        self.current_symbol = None
        self.last_close_prices = []

    def OnData(self, slice: Slice):
        if not slice.Bars.ContainsKey(self.qqq):
            return

        bar = slice[self.qqq]
        if self.window.IsReady:
            prev_close = self.window[0]
            log_return = np.log(bar.Close / prev_close)
            self.window.Add(bar.Close)
        else:
            if self.window.Count > 0:
                log_return = np.log(bar.Close / self.window[0])
            self.window.Add(bar.Close)

        # Track close prices for 10-day return
        self.last_close_prices.append(bar.Close)
        if len(self.last_close_prices) > self.holding_period:
            self.last_close_prices.pop(0)

    def OnWarmupFinished(self):
        self.Debug("Warmup complete.")
        self.current_symbol = self.qqq
        self.SetHoldings(self.current_symbol, 1)

    def DailyCheck(self):
        if not self.window.IsReady or len(self.last_close_prices) < self.holding_period:
            return

        # Compute volatility and VaR
        returns = [np.log(self.window[i] / self.window[i+1]) for i in range(self.lookback - 1)]
        vol = np.std(returns)
        var_10d = -self.z * vol * np.sqrt(self.holding_period)
        var_10d_up = self.z * vol * np.sqrt(self.holding_period)

        # Compute actual 10-day return
        recent_close = self.last_close_prices[-1]
        past_close = self.last_close_prices[0]
        actual_return = np.log(recent_close / past_close)

        # Switch logic
        if self.current_symbol == self.qqq and actual_return < var_10d:
            self.Debug(f"{self.Time.date()} VAR breach down → switching to TQQQ")
            self.SwitchTo(self.tqqq)

        elif self.current_symbol == self.tqqq and actual_return > var_10d_up:
            self.Debug(f"{self.Time.date()} VAR breach up → switching to QQQ")
            self.SwitchTo(self.qqq)

    def SwitchTo(self, symbol):
        self.Liquidate()
        self.SetHoldings(symbol, 1)
        self.current_symbol = symbol