"""
PRE-REGISTERED macro-state gates on the D1 4-index dip-z book (2026-10-08).

Written BEFORE the first run. Trial counter N = 6 (3 gates x 2 directions);
nothing here may be re-thresholded after seeing results -- any new variant
increments N and is reported as such.

Data (research/macro_data, free daily): CBOE VIX/VIX3M, FRED BAA10Y (Moody's
Baa-10Y spread; FRED's ICE HY OAS is truncated to 2023-10), index closes.
Causality: macro value of the SIGNAL bar's date (its close is known when the
bar closes); fill is the next open, as in backtest.simulate.

Gates (entry ALLOWED when true):
  G1a  VIX/VIX3M <= 1.0   (no backwardation)     G1b  VIX/VIX3M > 1.0
  G2a  VRP = VIX - RV10 > 5                        G2b  VRP <= 5
  G3a  BAA10Y 5d change <= 0 (credit not widening) G3b  5d change > 0
Hypotheses: block-stress (a) per the outside brief vs pay-for-stress (b) per
our own finding that expectancy rises with the vol regime. Both registered.

Survival rule (ALL required): vs the ungated book, ret/DD improves in 2008-09,
2022, IS (<=2016) and OOS (2017+); overall cadence stays >= 2/month; and
the gated book keeps >= 40% of trades.
"""
from __future__ import annotations
import sys, numpy as np, pandas as pd
sys.path.insert(0, ".")
import backtest as bt, portfolio as pf
from run_screen import zscore_entries

MD = "macro_data/"
IDX = pf.IDX
STOP = 3.0

def series():
    v = pd.read_csv(MD + "cboe_VIX.csv", parse_dates=["DATE"], date_format="%m/%d/%Y").set_index("DATE")["CLOSE"]
    v3 = pd.read_csv(MD + "cboe_VIX3M.csv", parse_dates=["DATE"], date_format="%m/%d/%Y").set_index("DATE")["CLOSE"]
    cr = pd.read_csv(MD + "BAA10Y.csv", index_col=0, parse_dates=True).iloc[:, 0].replace(".", np.nan).astype(float)
    return v, v3, cr

def gate_table(spx_close: pd.Series):
    v, v3, cr = series()
    df = pd.DataFrame({"vix": v, "vix3m": v3, "cr": cr}).sort_index().ffill()
    rv = np.log(spx_close / spx_close.shift(1)).rolling(10).std() * np.sqrt(252) * 100
    df["rv"] = rv.reindex(df.index).ffill()
    df["ratio"] = df.vix / df.vix3m
    df["vrp"] = df.vix - df.rv
    df["cr5"] = df.cr - df.cr.shift(5)
    return df

GATES = {
    "G1a": lambda g: g.ratio <= 1.0, "G1b": lambda g: g.ratio > 1.0,
    "G2a": lambda g: g.vrp > 5,      "G2b": lambda g: g.vrp <= 5,
    "G3a": lambda g: g.cr5 <= 0,     "G3b": lambda g: g.cr5 > 0,
}

def collect(gate=None, gt=None):
    out = []
    for s in IDX:
        d = bt.load(s, "PERIOD_D1")
        e, xma = zscore_entries(d, 20, -1.5, 0)
        if gate is not None:
            ok = gate(gt).reindex(pd.DatetimeIndex(d["ts"]).normalize()).to_numpy()
            ok = np.where(pd.isna(ok), False, ok).astype(bool)
            e = e & ok
        tr = bt.simulate(d, e, side=1, exit_ma=xma, max_bars=10, stop_atr=STOP)
        tr = bt.add_r(tr, d, stop_atr=STOP)
        for t in tr:
            t["symbol"] = s; t["exit_t"] = d["ts"][t["exit_i"]]
        out += tr
    return out

WINDOWS = {"2008-09": ("2008-01-01", "2009-04-01"), "2022": ("2022-01-01", "2023-01-01"),
           "IS<=2016": ("2008-01-01", "2017-01-01"), "OOS2017+": ("2017-01-01", "2030-01-01"),
           "FULL": ("2008-01-01", "2030-01-01")}

def window_stats(trades, lo, hi, risk=0.005):
    ts = [t for t in trades if np.datetime64(lo) <= t["t"] < np.datetime64(hi)]
    if len(ts) < 3: return None
    ts.sort(key=lambda t: t["exit_t"])
    eq = np.cumprod([1.0] + [1 + risk * t["r"] for t in ts]); pk = np.maximum.accumulate(eq)
    dd = ((pk - eq) / pk).max(); tot = eq[-1] - 1
    a = min(t["t"] for t in ts); b = max(t["exit_t"] for t in ts)
    mo = max((b - a) / np.timedelta64(365, "D") * 12, 1e-9)
    return dict(n=len(ts), pm=len(ts) / mo, bp=np.mean([t["ret"] for t in ts]) * 1e4,
                R=np.mean([t["r"] for t in ts]), tot=tot, dd=dd, rdd=(tot / dd if dd > 1e-9 else np.nan))

def main():
    spx = bt.load("SP500", "PERIOD_D1")
    gt = gate_table(pd.Series(spx["c"], index=pd.DatetimeIndex(spx["ts"]).normalize()))
    base = collect()
    runs = {"BASE": base}
    for k, f in GATES.items():
        runs[k] = collect(f, gt)
    print(f"{'run':<6}{'window':<10}{'n':>5}{'/mo':>6}{'bp':>7}{'meanR':>7}{'tot':>8}{'maxDD':>7}{'ret/DD':>8}")
    res = {}
    for k, tr in runs.items():
        for w, (lo, hi) in WINDOWS.items():
            s = window_stats(tr, lo, hi); res[(k, w)] = s
            if s: print(f"{k:<6}{w:<10}{s['n']:>5}{s['pm']:>6.1f}{s['bp']:>7.1f}{s['R']:>7.3f}{s['tot']:>8.1%}{s['dd']:>7.1%}{s['rdd']:>8.2f}")
            else: print(f"{k:<6}{w:<10}  <3 trades")
        print()
    print("SURVIVAL (rule: rdd up in all 4 sub-windows, cadence>=2/mo, >=40% of trades)")
    for k in GATES:
        ok = []
        for w in ("2008-09", "2022", "IS<=2016", "OOS2017+"):
            g, b = res[(k, w)], res[("BASE", w)]
            ok.append(bool(g and b and np.nan_to_num(g["rdd"], nan=-9) > np.nan_to_num(b["rdd"], nan=-9)))
        f = res[(k, "FULL")]
        keep = len(runs[k]) / len(base)
        card = f is not None and f["pm"] >= 2
        print(f"{k}: windows improved {ok}  cadence {f['pm'] if f else 0:.1f}/mo  kept {keep:.0%}  -> {'SURVIVES' if all(ok) and card and keep >= .4 else 'REJECTED'}")

if __name__ == "__main__":
    main()