"""Shielded stock sleeves: V9 risk machinery on top of stock selection.

The frontier battle so far lives in the leveraged-ETF trend family because
the cash-funded ledger forbids margin: only 3x instruments can reach the
CAGR bar. This lab quantifies the alternative the selection camp asks for —
take the BEST stock-selection sleeves (overnight_recovery top-10, low_risk
top-20) and wrap them in the V9/V10 machinery:

  regime shield : SPY > MA150 (through t-1) -> hold sleeve; else hold the
                  momentum leader of GLD/DBC
  vol target    : optional, scale = clip(target/vol60, min, 1), shift(1)
  trail stop    : optional basket-level trailing stop -> risk-off asset

PIT: all overlays use data through t-1 and one terminal shift, same
convention as TrendRiderV3.

    PYTHONPATH=. uv run python -m research.shield_lab
"""
from __future__ import annotations

import numpy as np
import pandas as pd

import trader
from strategies.base import BarStrategy
from strategies.robust_momentum import RobustMomentumStrategy

import research.tr9_lab as lab


class ShieldedMomentum(RobustMomentumStrategy):
    """Stock sleeve + SPY-regime shield (and optional vol target / trail)."""

    def __init__(self, signal: str, ma_long: int = 150, risk_off: tuple[str, ...] = ("GLD", "DBC"),
                 mom_lookback: int = 63, vol_target: float | None = None,
                 trail_pct: float | None = None, trail_window: int = 30, **kw) -> None:
        super().__init__(signal, **kw)
        self.ma_long = ma_long
        self.risk_off = risk_off
        self.mom_lookback = mom_lookback
        self.vol_target = vol_target
        self.trail_pct = trail_pct
        self.trail_window = trail_window

    def generate_weights(self, bars: dict[str, pd.DataFrame]) -> pd.DataFrame:
        w_s = super().generate_weights(bars)          # already .shift(1)'d, stocks only
        close = bars["close"]
        # keep the overlay instruments out of the base sleeve's selection
        extras_present = [c for c in ("SPY", *self.risk_off) if c in w_s.columns]
        if extras_present and w_s[extras_present].to_numpy().any():
            w_s = w_s.copy()
            w_s[extras_present] = 0.0
            w_s = w_s.div(w_s.sum(axis=1).replace(0.0, np.nan), axis=0).fillna(0.0)
        if "SPY" not in close.columns:
            return w_s
        spy = close["SPY"]
        risk_on = (spy > spy.rolling(self.ma_long).mean()).shift(1).fillna(False)

        # momentum leader among risk-off assets (through t-1)
        ro = [c for c in self.risk_off if c in close.columns]
        if ro:
            mom = close[ro].pct_change(self.mom_lookback).shift(1)
            leader = mom.idxmax(axis=1).where(mom.max(axis=1).notna(), ro[0])
        else:
            leader = pd.Series("", index=close.index)

        out = pd.DataFrame(0.0, index=close.index,
                           columns=sorted(set(w_s.columns) | set(ro)))
        out[w_s.columns] = w_s
        off_rows = ~risk_on
        for c in w_s.columns:
            out.loc[off_rows, c] = 0.0
        for i in np.where(off_rows)[0]:
            out.iat[i, out.columns.get_loc(leader.iloc[i])] = 1.0

        if self.trail_pct:
            # basket-level trailing stop on the sleeve NAV; ultra-conservative
            # lag: the trigger at output row i uses NAV through row i-2
            base_ret = (w_s * close[w_s.columns].pct_change(fill_method=None)).sum(axis=1)
            nav = (1 + base_ret.fillna(0.0)).cumprod()
            trig = ((nav / nav.rolling(self.trail_window, min_periods=5).max() - 1.0)
                    <= -self.trail_pct).shift(2).fillna(False)
            state, park_rows = False, []
            for i in range(len(trig)):
                if not risk_on.iloc[i]:
                    state = False
                elif trig.iloc[i]:
                    state = True
                park_rows.append(state)
            park = pd.Series(park_rows, index=close.index) & risk_on
            for i in np.where(park)[0]:
                for c in w_s.columns:
                    out.iat[i, out.columns.get_loc(c)] = 0.0
                out.iat[i, out.columns.get_loc(leader.iloc[i])] = 1.0

        if self.vol_target:
            rets = (out.shift(0) * close.pct_change(fill_method=None).reindex_like(out)).sum(axis=1)
            realized = rets.rolling(60, min_periods=21).std() * np.sqrt(252)
            scale = (self.vol_target / realized).clip(0.6, 1.0).shift(1).fillna(1.0)
            out = out.mul(scale, axis=0)
        return out


def register() -> list[str]:
    reg = trader.STRATEGY_REGISTRY
    names = []
    variants = {
        "sh_orr_shield": lambda: ShieldedMomentum("overnight_recovery", top_n=10, rebal_freq=10,
                                                  buffer=2.0),
        "sh_orr_shield_vt": lambda: ShieldedMomentum("overnight_recovery", top_n=10, rebal_freq=10,
                                                     buffer=2.0, vol_target=0.28),
        "sh_orr_full": lambda: ShieldedMomentum("overnight_recovery", top_n=10, rebal_freq=10,
                                                buffer=2.0, vol_target=0.28, trail_pct=0.20),
        "sh_lowrisk_shield_vt": lambda: ShieldedMomentum("low_risk", top_n=20, rebal_freq=21,
                                                         buffer=2.0, vol_target=0.22),
    }
    for n, f in variants.items():
        reg[n] = f
        trader.MIXED_STRATEGY_EXTRA_TICKERS[n] = ["SPY", "GLD", "DBC"]
        names.append(n)
    return names


def main() -> None:
    names = register()
    curves = lab.run(["overnight_recovery_top10", "lowrisk_top20", "us_three_sleeve",
                      "trend_rider_v10"] + names, end="2026-10-07")
    lab.report(curves, "shielded sleeves vs frontier")


if __name__ == "__main__":
    main()