from AlgorithmImports import *


class VIXDualStrategy(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2016, 1, 1)
        self.set_end_date(2026, 7, 22)
        self.set_cash(1_000_000)
        # Add SPY as an asset to trade.
        self._spy = self.add_equity("SPY")
        # Add some indicators to calculate trailing returns.
        self._spy.daily_returns = self.roc(self._spy, 1, Resolution.DAILY)
        self._spy.daily_returns.window[8]
        self._spy.previous_close = self.identity(self._spy, Resolution.DAILY)
        # Add the VIX and VIX3M Indices.
        self._vix = self.add_index("VIX")
        self._vix3m = self.add_index("VIX3M")
        # Add the VIXY, which tracks the S&P 500 VIX Short-Term Futures Index.
        self._vixy = self.add_equity("VIXY")
        # Add a warm-up period to prime SPY's trailing daily returns.
        self.set_warm_up(10, Resolution.DAILY)    
        # Add a Scheduled Event to rebalance the portfolio each day.
        # 16 minutes before the close is the last chance to place MOC orders.
        self.schedule.on(self.date_rules.every_day(self._spy), self.time_rules.before_market_close(self._spy, 16), self._rebalance)

    def _rebalance(self) -> None:
        # During warm-up, do nothing.
        if self.is_warming_up:
            return
        # Get the trailing returns over the last 10 days.
        # Use the return from the previous close to now as the latest "daily return".
        returns = [x.value for x in self._spy.daily_returns.window]
        returns.append(self._spy.price / self._spy.previous_close.current.value - 1)
        # Calculate the expected realized vol (annualised, VIX points).
        e_rv30 = np.std(returns, ddof=1) * np.sqrt(252) * 100
        # Calculate the expected VRP.
        e_vrp = self._vix.price - e_rv30
        # Determine the target weight for VIXY.
        if e_vrp > 0 and self._vix.price < self._vix3m.price:
            # Case 1: Full short-vol conviction
            target_weight = -self._vix.price / 100
        elif e_vrp < 0 and self._vix.price < self._vix3m.price:
            # Case 2: Medium short-vol conviction (half size)
            target_weight = -0.5 * self._vix.price / 100
        elif e_vrp < 0 and self._vix.price > self._vix3m.price:
            # Case 3: Full long-vol conviction
            target_weight = self._vix.price / 100
        else:
            # Case 4: Cash (conflicting signals or eVRP == 0)
            target_weight = 0
        # To reduce churn, only rebalance when the target weight changes sign or is at least 
        # 2% away from the current weight.
        current_weight = self._vixy.holdings.holdings_value / self.portfolio.total_portfolio_value
        if np.sign(self._vixy.holdings.holdings_value) == np.sign(target_weight) and abs(current_weight - target_weight) <= 0.02:
            return
        # Rebalance VIXY to the target weight.
        vixy_qty = self.calculate_order_quantity(self._vixy, target_weight)
        if vixy_qty:
            self.market_on_close_order(self._vixy, vixy_qty)
        # If we are short volaility, buy SPY to keep gross exposure at 100%
        # instead of just holding onto cash.
        if target_weight < 0:
            spy_qty = self.calculate_order_quantity(self._spy, 1+target_weight)
            if spy_qty:
                self.market_on_close_order(self._spy, spy_qty)
        # Otherwise, exit SPY.
        elif self._spy.invested:
            self.market_on_close_order(self._spy, -self._spy.holdings.quantity)