"""
Left-tail control on the S&P 500 ∪ S&P 400 union book: does a daily
market-regime shield change the universe_midcap verdict?

Background (research/lab/universe_midcap.py): ranking S&P 400 names inside
overnight_recovery_top10 cost 4.7pp of CAGR and 0.12 of Sharpe (2019-26),
and its excess losses concentrated in down years (2022: -26pp). This applies
the same overlay to both arms and asks whether the union becomes acceptable
once the crash tail is cut.

    PYTHONPATH=. uv run python -m research.lab.union_shield

Prereg (written before the run): exposure_t = 1{SPY close through t-1 above
its 150-day average through t-1}, else cash (0 return); applied to the
tranche-averaged daily book returns saved by universe_midcap.py;
20 bp charged per unit of exposure change. One mechanism, no variants.

Result (2026-10-09): FAIL. The shield hurts every block of the base book
(whipsaw: it sits out 24-38% of 2020/2021/2025 bull legs) — base
29.4%/0.98 -> base+shield 16.1%/0.70 — and union+shield (15.5%/0.67) lands
below both base+shield (P1) and the plain unshielded base (P2). MaxDDs are
unchanged to the decimal (-34.3% base, -38.8% union+shield), so the filter
buys no tail either. The union ordering (union >= adv20 > minus_top3) is
unchanged under the shield. Conclusion: left-tail control does not salvage
the S&P 400 extension; a SPY>MA150 regime filter is simply bad for this
momentum book.
"""

from __future__ import annotations

import json
import os
import warnings

import numpy as np
import pandas as pd

from research.lab.alpha import alpha_stats
from research.lab.panel_mid import load_us_mid

warnings.filterwarnings("ignore")
RESULTS = os.path.join("data", "research", "results", "universe_midcap")
OUT = os.path.join("data", "research", "results", "union_shield")
PRIMARY = ("2019-01-02", "2026-09-24")
BLOCKS = [("2019-01-02", "2020-12-31"), ("2021-01-04", "2022-12-30"),
          ("2023-01-03", "2024-12-31"), ("2025-01-02", "2026-09-24")]
REFERENCE = ("2007-01-03", "2018-12-31")
COST = 0.002
PREREG = {
    "question": "With a SPY>MA150 regime shield (else cash) applied to both arms, does adding PIT "
                "S&P 400 members to overnight_recovery_top10 still lose?",
    "arms": ["base", "base+shield", "union", "union+shield", "union_adv20+shield (reported)"],
    "mechanics": "exposure_t = 1{SPY close through t-1 > its 150d average through t-1}, else 0 (cash); "
                 "r'_t = e_t r_t - 0.002 |e_t - e_{t-1}|; inputs = universe_midcap/returns.pkl "
                 "(tranche-averaged, 20 bp a side already inside)",
    "primary_window": PRIMARY,
    "pass_if_all": [
        "P1: Sharpe(union+shield) > Sharpe(base+shield)  (universe A/B under the same overlay)",
        "P2: Sharpe(union+shield) >= Sharpe(base)  (vs the live book, unshielded)",
        "P3: MaxDD(union+shield) >= MaxDD(base) - 5pp",
    ],
    "on_pass": "the union deserves a shield-equipped forward check; no production change from this alone",
    "status": "the window was already examined; not independent OOS; one mechanism, no variants",
}


def window_stats(r: pd.Series, a: str, b: str) -> dict:
    x = r.loc[a:b]
    e = (1 + x).cumprod()
    yrs = len(x) / 252
    return {"cagr": float(e.iloc[-1] ** (1 / yrs) - 1),
            "sharpe": float(x.mean() / x.std() * np.sqrt(252)),
            "mdd": float((e / e.cummax() - 1).min())}


def main() -> None:
    os.makedirs(OUT, exist_ok=True)
    res_path = os.path.join(OUT, "results.json")
    if os.path.exists(res_path):
        raise SystemExit(f"{res_path} exists; this study does not overwrite itself")
    if not os.path.exists(os.path.join(OUT, "prereg.json")):
        with open(os.path.join(OUT, "prereg.json"), "w") as fh:
            json.dump(PREREG, fh, indent=2)

    R = pd.read_pickle(os.path.join(RESULTS, "returns.pkl"))
    p = load_us_mid("mid")
    spy = p.extra["etf"]["SPY"]
    ma = spy.rolling(150, min_periods=120).mean()
    expo = (spy.shift(1) > ma.shift(1)).astype(float).reindex(R.index).ffill().fillna(0.0)
    d = expo.diff().abs().fillna(0.0)

    S = {}
    for arm in R.columns:
        S[arm] = R[arm]
        S[f"{arm}+shield"] = expo * R[arm] - COST * d

    res = {}
    for arm, r in S.items():
        res[arm] = {w: window_stats(r, *ab) for w, ab in
                    [("primary", PRIMARY), ("reference", REFERENCE)]}
        res[arm]["blocks"] = [round(window_stats(r, *ab)["sharpe"], 3) for ab in BLOCKS]
        res[arm]["years"] = {int(k): round(float(v), 4) for k, v in
                             ((1 + r.loc[PRIMARY[0]:PRIMARY[1]]).groupby(
                                 r.loc[PRIMARY[0]:PRIMARY[1]].index.year).prod() - 1).items()}

    a = alpha_stats(S["union+shield"].loc[PRIMARY[0]:PRIMARY[1]],
                    S["base"].loc[PRIMARY[0]:PRIMARY[1]])
    res["union+shield"]["alpha_vs_base"] = {k: float(a[k]) for k in ("alpha", "t")}
    res["exposure_share_by_year"] = {int(y): round(float(v), 3) for y, v in
                                     expo.loc[PRIMARY[0]:PRIMARY[1]].groupby(
                                         expo.loc[PRIMARY[0]:PRIMARY[1]].index.year).mean().items()}

    base_, base_sh = res["base"]["primary"], res["base+shield"]["primary"]
    union_sh = res["union+shield"]["primary"]
    checks = {"P1_union_shield_gt_base_shield": union_sh["sharpe"] > base_sh["sharpe"],
              "P2_union_shield_ge_base": union_sh["sharpe"] >= base_["sharpe"],
              "P3_mdd_not_worse_5pp": union_sh["mdd"] >= base_["mdd"] - 0.05}
    verdict = "PASS" if all(checks.values()) else "FAIL"
    summary = {"prereg": PREREG, "results": res, "checks": {k: bool(v) for k, v in checks.items()},
               "verdict": verdict}
    with open(res_path, "w") as fh:
        json.dump(summary, fh, indent=2, default=float)
    pd.DataFrame(S).to_pickle(os.path.join(OUT, "returns.pkl"))

    for arm in ("base", "base+shield", "union", "union+shield", "union_adv20+shield", "union_minus_top3+shield"):
        s = res[arm]
        print(f"{arm:22s} primary: cagr {s['primary']['cagr']*100:5.1f}% sh {s['primary']['sharpe']:.2f} "
              f"mdd {s['primary']['mdd']*100:5.1f}%  blocks {s['blocks']}  ref {s['reference']['sharpe']:.2f}")
    print("exposure share by year:", res["exposure_share_by_year"])
    print("union+shield alpha vs base:", res["union+shield"]["alpha_vs_base"])
    print("checks:", summary["checks"], "→", verdict)

    from research.lab.trials import append_trials
    append_trials(os.path.join("data", "research", "results", "trials_us.jsonl"),
                  [{"name": f"union_shield:{k}", **{f"{w}_{x}": y for w in ("primary", "reference")
                    for x, y in v.items() if isinstance(y, (int, float))}}
                   for k, v in res.items() if isinstance(v, dict) and "primary" in v],
                  {"study": "union_shield", "role": "preregistered overlay on a previously seen window"})


if __name__ == "__main__":
    main()