"""Short-volume veto overlay test on the momentum sleeve (official engine).

Hypothesis from the first scan (research/lab/short_volume.py): names in the
top quintile of 21-day mean REG SHO short-volume ratio (SVR = ShortVolume /
TotalVolume) underperform over the next month (S&P 500: IC21 -0.020, t -3.1;
Q5-Q1 spread -0.40%/mo). Test: veto those names at selection time (score ->
NaN before top-N selection), on the flagship overnight_recovery_top10 sleeve.

PIT: the SVR panel is shifted by 2 sessions inside the veto (REG SHO files
are published after the close), matching the conservative convention of the
scan.

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

import pickle

import numpy as np
import pandas as pd

import trader
from strategies.robust_momentum import RobustMomentumStrategy, select_with_buffer

import research.tr9_lab as lab

SVR_PATH = "data/research/short_volume.pkl"
_SVR_CACHE: dict = {}


def svr21_panel() -> pd.DataFrame:
    if "svr" not in _SVR_CACHE:
        with open(SVR_PATH, "rb") as fh:
            d = pickle.load(fh)
        svr = (d["short"] / d["total"].replace(0.0, np.nan))
        _SVR_CACHE["svr"] = svr.rolling(21, min_periods=10).mean()
    return _SVR_CACHE["svr"]


class RobustMomentumSvr(RobustMomentumStrategy):
    """RobustMomentum + veto of top-quantile short-volume names at selection."""

    def __init__(self, signal: str, svr_quant: float = 0.8, svr_lag: int = 2, **kw) -> None:
        super().__init__(signal, **kw)
        self.svr_quant = svr_quant
        self.svr_lag = svr_lag

    def generate_weights(self, bars: dict[str, pd.DataFrame]) -> pd.DataFrame:
        close = bars["close"]
        panel = {f: bars[f].reindex(index=close.index, columns=close.columns) for f in self.fields}
        if "eligible" in bars:
            panel["eligible"] = bars["eligible"]
        score = self._score(panel)
        eligible = close.notna()
        if "eligible" in bars:
            eligible &= bars["eligible"].reindex_like(close).fillna(False).astype(bool)
        score = score.where(eligible)

        svr = svr21_panel().reindex(index=close.index, columns=close.columns).shift(self.svr_lag)
        q = svr.rank(axis=1, pct=True)
        veto = (q >= self.svr_quant) & svr.notna()
        score = score.where(~veto)

        weights = select_with_buffer(score, self.top_n, self.rebal_freq, self.buffer, self.warmup,
                                     prices=close)
        return weights.shift(1).fillna(0.0)


def register() -> list[str]:
    reg = trader.STRATEGY_REGISTRY
    names = []
    reg["overnight_recovery_svr80"] = lambda **kw: RobustMomentumSvr(
        "overnight_recovery", top_n=10, rebal_freq=10, buffer=2.0, svr_quant=0.8)
    names.append("overnight_recovery_svr80")
    reg["overnight_recovery_svr90"] = lambda **kw: RobustMomentumSvr(
        "overnight_recovery", top_n=10, rebal_freq=10, buffer=2.0, svr_quant=0.9)
    names.append("overnight_recovery_svr90")
    reg["overnight_recovery_svr70"] = lambda **kw: RobustMomentumSvr(
        "overnight_recovery", top_n=10, rebal_freq=10, buffer=2.0, svr_quant=0.7)
    names.append("overnight_recovery_svr70")
    return names


def main() -> None:
    names = register()
    curves = lab.run(["overnight_recovery_top10", "lowrisk_top20"] + names)
    lab.report(curves, "SVR veto overlay")


if __name__ == "__main__":
    main()