# region imports
from AlgorithmImports import *

import numpy as np
# endregion

# Implementation of Strategy 4 ("eVRP + BoC + Sizing") from:
#   Zarattini, C., Mele, A., & Aziz, A. (2025). "The Volatility Edge: A Dual Approach
#   for VIX ETNs Trading." Swiss Finance Institute.
#
# Signal (paper, Table 2):
#   eVRP > 0  and VIX < VIX3M -> short vol at VIX%
#   eVRP <= 0 and VIX < VIX3M -> short vol at 0.5 x VIX%
#   eVRP <= 0 and VIX > VIX3M -> long vol at VIX%
#   otherwise                 -> cash
# where eVRP = VIX - std(last 10 SPY daily returns) * sqrt(252) * 100, the final return
# running from yesterday's close to the current intraday price.
#
# Deviations from the paper, and why:
#   - Default reality models only. The paper nets a flat 5 bps per trade; overriding
#     LEAN's fee model to reproduce that is not permitted, so absolute returns here will
#     differ from Table 3.
#   - Signal fires 16 minutes before the close rather than the paper's 15. LEAN's
#     default market-on-close submission buffer is 15m30s and rejects anything later.
#   - Backtest covers the trailing 5 years, so the paper's 2008 (+86.9%) and 2020
#     (+42.3%) years are outside the sample.
#   - The paper trades synthetic proxies: VIXLONG (SPVXSTR less 50bp/yr) and VIXSHORT
#     (the inverse index, XIV-like -1x, less 80bp). "Long VIXSHORT at VIX%" is -VIX% of
#     1x futures exposure, so any listed expression with matching exposure is
#     equivalent. Legs are therefore defined as (ticker, beta): short VIXY at 1x, long
#     SVXY at 2x, or long SVIX at 1x.
#   - The paper's Strategy 4 rule box writes "eVRP < 0" while its Table 2 summary writes
#     "eVRP <= 0" for the same strategy. Table 2 and the reference notebook agree, so
#     <= is used here.


class VIXDualStrategy(QCAlgorithm):
    """Dual-signal VIX ETN strategy with VIX-scaled position sizing."""

    def initialize(self) -> None:
        # self.set_end_date(2026, 7, 22)
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.initial_cash = 100000
        self.set_cash(self.initial_cash)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0

        lookback = int(self.get_parameter("lookback", 10))
        minutes_before_close = int(self.get_parameter("minutes_before_close", 16))
        maximum_exposure = float(self.get_parameter("maximum_exposure", 1.0))
        sizing_divisor = float(self.get_parameter("sizing_divisor", 100.0))
        # "dynamic" = Strategy 4 (eVRP + BoC + Sizing); "static" = Strategy 3 (eVRP + BoC).
        sizing_mode = str(self.get_parameter("sizing_mode", "dynamic"))
        static_size = float(self.get_parameter("static_size", 0.20))
        # Band is quoted in vol-exposure terms and converted per instrument below.
        exposure_band = float(self.get_parameter("exposure_band", 0.02))
        # The paper holds unused capital in cash and assumes no interest on it. Its SPY
        # blend is a separate month-end-rebalanced sleeve, not a daily gross-100% overlay.
        spy_overlay = bool(int(self.get_parameter("spy_overlay", 1)))

        # Leg definitions. Beta is exposure to 1x short-term VIX futures: VIXY/VXX = +1.0,
        # SVXY = -0.5 (it was -1.0 before Feb 2018), SVIX = -1.0.
        long_ticker = self.get_parameter("long_vol_etn", "VIXY")
        long_beta = float(self.get_parameter("long_vol_beta", 1.0))
        # Leave short_vol_etn equal to long_vol_etn to express short vol by shorting the
        # +1x ETN; set it to SVXY or SVIX to go long an inverse ETN instead.
        short_ticker = self.get_parameter("short_vol_etn", long_ticker)
        default_short_beta = 1.0 if short_ticker == long_ticker else -0.5
        short_beta = float(self.get_parameter("short_vol_beta", default_short_beta))

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        # Signal inputs. SPY is a data feed unless the overlay parameter is switched on.
        # It doubles as the benchmark, so it is NOT added a second time at daily
        # resolution: the signal needs the minute bar at 3:44pm to build today's return.
        self.benchmark_symbol = "SPY"
        self.spy_prices = []
        equity = self.add_equity(self.benchmark_symbol, Resolution.MINUTE)
        self.set_benchmark(equity.symbol)
        spy = equity.symbol
        vix = self.add_index("VIX", Resolution.MINUTE).symbol
        vix3m = self.add_index("VIX3M", Resolution.MINUTE).symbol

        long_symbol = Symbol.create(long_ticker, SecurityType.EQUITY, Market.USA)
        short_symbol = Symbol.create(short_ticker, SecurityType.EQUITY, Market.USA)
        tradables = [long_symbol] if short_symbol == long_symbol else [long_symbol, short_symbol]

        self.universe_settings.resolution = Resolution.MINUTE
        self.set_universe_selection(ManualUniverseSelectionModel(tradables))

        self.set_alpha(
            VIXDualSignalAlphaModel(
                (long_symbol, long_beta),
                (short_symbol, short_beta),
                spy,
                vix,
                vix3m,
                lookback,
                minutes_before_close,
                sizing_divisor,
                sizing_mode,
                static_size,
                spy_overlay,
            )
        )

        # A leg held at 1/|beta| notional needs a 1/|beta| wider notional band to keep the
        # same exposure band, otherwise an inverse leg churns twice as often.
        bands = {
            long_symbol: exposure_band / abs(long_beta),
            short_symbol: exposure_band / abs(short_beta),
            spy: exposure_band / abs(long_beta),
        }
        self.set_portfolio_construction(BandedWeightPortfolioConstructionModel(bands, exposure_band))
        self.set_execution(MarketOnCloseExecutionModel())
        self.set_risk_management(MaximumGrossExposureRiskModel(maximum_exposure))

        # Warm up historical data before placing initial trades. The 5-day boilerplate
        # floor is not enough here: 10 daily returns need 11 daily bars, and a short
        # warm-up would silently compute the vol estimate on fewer samples.
        self.set_warm_up(lookback + 1, Resolution.DAILY)


class VIXDualSignalAlphaModel(AlphaModel):
    """Generates the dual (eVRP + term-structure) volatility signal once per day.

    The signal is produced as a vol exposure in units of 1x short-term VIX futures. Each
    leg converts that exposure into a notional weight by dividing by the instrument's
    beta, so long SVXY (beta -0.5) is sized at twice the notional of a short VIXY
    position (beta +1.0) for identical exposure.
    """

    def __init__(
        self,
        long_vol: Tuple[Symbol, float],
        short_vol: Tuple[Symbol, float],
        spy: Symbol,
        vix: Symbol,
        vix3m: Symbol,
        lookback: int = 10,
        minutes_before_close: int = 16,
        sizing_divisor: float = 100.0,
        sizing_mode: str = "dynamic",
        static_size: float = 0.20,
        spy_overlay: bool = False,
    ) -> None:
        self.name = "VIXDualSignalAlphaModel"
        self._long_vol_symbol, self._long_vol_beta = long_vol
        self._short_vol_symbol, self._short_vol_beta = short_vol
        self._spy = spy
        self._vix = vix
        self._vix3m = vix3m
        self._lookback = lookback
        self._minutes_before_close = minutes_before_close
        self._sizing_divisor = sizing_divisor
        self._sizing_mode = sizing_mode.lower()
        self._static_size = static_size
        self._spy_overlay = spy_overlay

        self._daily_returns: RateOfChange = None
        self._previous_close: Identity = None
        self._last_signal_date: date = None

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        self._ensure_indicators(algorithm)
        if algorithm.is_warming_up or not self._is_signal_time(algorithm):
            return []

        exposure = self._vol_exposure(algorithm)
        self._last_signal_date = algorithm.time.date()
        if exposure is None:
            return []

        # Route the exposure to whichever leg expresses it, and flatten the other.
        separate_short_leg = self._short_vol_symbol != self._long_vol_symbol
        if exposure < 0 and separate_short_leg:
            long_weight = 0.0
            short_weight = exposure / self._short_vol_beta
        else:
            long_weight = exposure / self._long_vol_beta
            short_weight = 0.0

        # If the leg carrying the exposure has no price, the vol trade cannot be placed.
        # Emitting the SPY overlay anyway would leave a naked long-equity book
        # masquerading as a vol strategy, which is what happens when the backtest starts
        # before the ETN's inception date.
        carrier = self._short_vol_symbol if (exposure < 0 and separate_short_leg) else self._long_vol_symbol
        if exposure != 0 and algorithm.securities[carrier].price == 0:
            algorithm.log(f"{carrier.value} has no price; suppressing insights (before inception?)")
            return []

        insights = [self._make_insight(self._long_vol_symbol, long_weight)]
        if separate_short_leg:
            insights.append(self._make_insight(self._short_vol_symbol, short_weight))

        # Optional overlay: park the capital not used by the short-vol leg in SPY so gross
        # exposure stays at 100%. Not in the paper; it stacks two risks that draw down
        # together. Off by default.
        if self._spy_overlay:
            used = abs(long_weight) + abs(short_weight)
            spy_weight = max(0.0, 1 - used) if exposure < 0 else 0.0
            insights.append(self._make_insight(self._spy, spy_weight))

        return insights

    def _ensure_indicators(self, algorithm: QCAlgorithm) -> None:
        """Register the SPY indicators on first use."""
        if self._daily_returns is not None:
            return
        self._daily_returns = algorithm.roc(self._spy, 1, Resolution.DAILY)
        # The window holds the last (lookback - 1) completed daily returns; the final
        # return of the sample is today's partial close-to-now return.
        self._daily_returns.window.size = self._lookback - 1
        self._previous_close = algorithm.identity(self._spy, Resolution.DAILY)

    def _is_signal_time(self, algorithm: QCAlgorithm) -> bool:
        """True once per trading day, `minutes_before_close` before the SPY close.

        Derived from exchange hours rather than a fixed clock time so that half-days
        (1:00pm closes) are handled correctly, per the paper's footnote 18.
        """
        if self._last_signal_date == algorithm.time.date():
            return False
        exchange = algorithm.securities[self._spy].exchange
        if not exchange.date_time_is_open(algorithm.time):
            return False
        market_close = exchange.hours.get_next_market_close(algorithm.time, False)
        seconds_to_close = (market_close - algorithm.time).total_seconds()
        return 0 < seconds_to_close <= self._minutes_before_close * 60

    def _vol_exposure(self, algorithm: QCAlgorithm) -> float:
        """Signed target exposure in 1x VIX-futures units, or None if inputs are stale."""
        if self._daily_returns.window.count < self._lookback - 1:
            return None
        if not self._previous_close.is_ready or self._previous_close.current.value == 0:
            return None

        spy_price = algorithm.securities[self._spy].price
        vix = algorithm.securities[self._vix].price
        vix3m = algorithm.securities[self._vix3m].price
        if spy_price == 0 or vix == 0 or vix3m == 0:
            return None

        # Trailing daily returns, with today's partial return as the most recent sample.
        returns = [x.value for x in self._daily_returns.window]
        returns.append(spy_price / self._previous_close.current.value - 1)

        # Expected realised vol over the next 30 days, annualised in VIX points.
        e_rv30 = float(np.std(returns, ddof=1) * np.sqrt(252) * 100)
        e_vrp = vix - e_rv30
        contango = vix < vix3m
        backwardation = vix > vix3m

        # Strategy 4 scales with the VIX level; Strategy 3 uses a fixed 20%/10%. The
        # signal logic is identical either way, only the size responds.
        base = vix / self._sizing_divisor if self._sizing_mode == "dynamic" else self._static_size

        if e_vrp > 0 and contango:
            exposure = -base          # Case 1: full short-vol conviction
        elif e_vrp <= 0 and contango:
            exposure = -0.5 * base    # Case 2: half-size short-vol conviction
        elif e_vrp <= 0 and backwardation:
            exposure = base           # Case 3: full long-vol conviction
        else:
            exposure = 0.0            # Case 4: conflicting signals -> cash

        algorithm.log(
            f"eRV30={e_rv30:.2f} VIX={vix:.2f} VIX3M={vix3m:.2f} "
            f"eVRP={e_vrp:.2f} mode={self._sizing_mode} exposure={exposure:.4f}"
        )
        return exposure

    def _make_insight(self, symbol: Symbol, weight: float) -> Insight:
        if weight > 0:
            direction = InsightDirection.UP
        elif weight < 0:
            direction = InsightDirection.DOWN
        else:
            direction = InsightDirection.FLAT

        # Sign travels in the direction, magnitude in the weight.
        return Insight.price(symbol, timedelta(days=1), direction, None, None, self.name, abs(weight))


class BandedWeightPortfolioConstructionModel(PortfolioConstructionModel):
    """Turns a signed insight weight into an absolute share target, with a churn band.

    A new target is only issued when the target weight changes sign or moves further than
    that symbol's band away from the current weight. Bands are supplied per symbol so a
    2x-sized inverse leg is not rebalanced twice as often as a 1x leg.
    """

    def __init__(self, bands: Dict[Symbol, float], default_band: float = 0.02) -> None:
        self._bands = bands
        self._default_band = default_band

    def create_targets(self, algorithm: QCAlgorithm, insights: List[Insight]) -> List[PortfolioTarget]:
        targets: List[PortfolioTarget] = []
        total_value = algorithm.portfolio.total_portfolio_value
        if total_value <= 0:
            return targets

        for insight in insights:
            symbol = insight.symbol
            weight = insight.weight if insight.weight is not None else 0.0
            target_weight = self._sign_of(insight.direction) * float(weight)
            current_weight = algorithm.portfolio[symbol].holdings_value / total_value
            band = self._bands.get(symbol, self._default_band)

            # Skip only when we are already on the right side and close enough.
            same_side = np.sign(current_weight) == np.sign(target_weight)
            if same_side and abs(current_weight - target_weight) <= band:
                continue

            target = PortfolioTarget.percent(algorithm, symbol, target_weight)
            if target is None:
                continue
            targets.append(target)
        return targets

    @staticmethod
    def _sign_of(direction: InsightDirection) -> float:
        """InsightDirection is a .NET enum and is not castable to float in Python."""
        if direction == InsightDirection.UP:
            return 1.0
        if direction == InsightDirection.DOWN:
            return -1.0
        return 0.0


class MarketOnCloseExecutionModel(ExecutionModel):
    """Fills portfolio targets with MOC orders, netting against existing holdings."""

    def execute(self, algorithm: QCAlgorithm, targets: List[PortfolioTarget]) -> None:
        if algorithm.is_warming_up:
            return
        for target in targets:
            symbol = target.symbol
            open_quantity = sum(
                ticket.quantity
                for ticket in algorithm.transactions.get_open_order_tickets(lambda t, s=symbol: t.symbol == s)
            )
            delta = target.quantity - algorithm.portfolio[symbol].quantity - open_quantity
            if delta != 0 and algorithm.securities[symbol].is_tradable:
                algorithm.market_on_close_order(symbol, delta)


class MaximumGrossExposureRiskModel(RiskManagementModel):
    """Hard cap on absolute exposure to any single security."""

    def __init__(self, maximum_exposure: float = 1.0) -> None:
        self._maximum_exposure = abs(maximum_exposure)

    def manage_risk(self, algorithm: QCAlgorithm, targets: List[PortfolioTarget]) -> List[PortfolioTarget]:
        risk_targets: List[PortfolioTarget] = []
        total_value = algorithm.portfolio.total_portfolio_value
        if total_value <= 0:
            return risk_targets

        for kvp in algorithm.portfolio:
            holding = kvp.Value
            if not holding.invested:
                continue
            weight = holding.holdings_value / total_value
            if abs(weight) > self._maximum_exposure:
                capped = np.sign(weight) * self._maximum_exposure
                algorithm.debug(f"Risk cap hit on {holding.symbol}: {weight:.2f} -> {capped:.2f}")
                risk_targets.append(PortfolioTarget.percent(algorithm, holding.symbol, capped))
        return risk_targets