#!/usr/bin/env python3
"""Full gauntlet on the equity low-vol/lottery-aversion lead (liquid high-vol cohort).

The make-or-break for the first non-crypto edge: (1) name-bootstrap (robust to WHICH names? —
the test that killed carry), (2) cohort-param robustness (did $10M/60pct manufacture it?),
(3) realistic short-borrow + long-only (is the edge trapped in the un-cheap short?),
(4) honest deflation + per-year, (5) correlation to the crypto momentum book (diversifier?).
Free — DBEQ on disk.
"""
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from equity_factor_gate import load, roll  # noqa: E402
from signal_sweep import xs_weights, validate, sharpe_t  # noqa: E402
import pit_sweep  # noqa: E402
from surfer_poc import compute_weights, CFG  # noqa: E402
import torch  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"


def main():
    insts, days, close, dvol = load()
    T, N = close.shape
    lc = np.log(close)
    R = np.zeros((T, N)); R[1:] = lc[1:] - lc[:-1]; R = np.where(np.isfinite(R), R, 0.0)
    dv30 = roll(np.mean, np.nan_to_num(dvol), 30)
    vol63 = roll(np.std, R, 63)
    year = (1970 + days / 365.25).astype(int)
    T_ = lambda x: torch.tensor(x[np.isfinite(x)], device=DEV, dtype=torch.float64)

    def cohort(dv_floor, vol_pct):
        u = np.zeros((T, N), bool)
        for t in range(T):
            liq = (dv30[t] > dv_floor) & np.isfinite(close[t]) & np.isfinite(vol63[t])
            e = np.where(liq)[0]
            if len(e) > 50:
                u[t, e[vol63[t, e] >= np.quantile(vol63[t, e], vol_pct)]] = True
        return u

    def smooth_weekly(w, K=5):
        a = 2.0 / 6
        for t in range(1, T):
            w[t] = a * w[t] + (1 - a) * w[t - 1]
        wh = w.copy(); last = 0
        for t in range(T):
            if t % K == 0:
                last = t
            wh[t] = w[last]
        return wh

    def lowvol_pnl(univ, cols=None, long_only=False, borrow_ann=0.0):
        sig = (-vol63).copy(); sig[~univ] = np.nan
        if cols is not None:
            mask = np.zeros(N, bool); mask[cols] = True; sig[:, ~mask] = np.nan
        if long_only:
            w = np.zeros((T, N))
            for t in range(T):
                e = np.where(np.isfinite(sig[t]))[0]
                if len(e) > 20:
                    k = max(int(0.10 * len(e)), 5); w[t, e[np.argsort(-sig[t, e])[:k]]] = 1.0 / k
        else:
            w = xs_weights(sig)
        w = smooth_weekly(w)
        g = np.sum(w[:-1] * R[1:], axis=1)
        g -= np.sum(np.abs(w[1:] - w[:-1]), axis=1) * 10.0 / 1e4              # 10bp trade cost
        if borrow_ann > 0:                                                    # borrow on short notional
            short_notional = np.sum(np.clip(-w[:-1], 0, None), axis=1)
            g -= short_notional * borrow_ann / 252.0
        return g

    base = cohort(1e7, 0.60)
    base_pnl = lowvol_pnl(base)
    v = validate(base_pnl, days, 30)
    print(f"\n===== EQUITY LOW-VOL GAUNTLET (liquid high-vol cohort, deflate N=30) =====")
    print(f"BASE L/S: full {v['full']:+.2f} OOS {v['oos']:+.2f} CPCVmed {v['med']:+.2f} DSR {v['dsr']:.2f}")

    # (1) name-bootstrap
    cohort_cols = np.where(base.any(0))[0]
    rng = np.random.default_rng(5); sh = []
    for _ in range(200):
        c = rng.choice(cohort_cols, size=max(len(cohort_cols) // 2, 20), replace=False)
        s = sharpe_t(T_(lowvol_pnl(base, cols=c)))
        if not math.isnan(s):
            sh.append(s)
    sh = np.array(sh)
    print(f"(1) name-bootstrap(200): frac>0 {np.mean(sh>0):.2f}  median {np.median(sh):+.2f}  5th {np.percentile(sh,5):+.2f}")

    # (2) cohort-param robustness
    print("(2) cohort-param grid (L/S full Sharpe):")
    for fl in [5e6, 1e7, 2e7, 5e7]:
        row = []
        for vp in [0.50, 0.60, 0.70]:
            row.append(sharpe_t(T_(lowvol_pnl(cohort(fl, vp)))))
        print(f"     ${fl/1e6:>3.0f}M floor: " + "  ".join(f"vp{int(vp*100)}:{r:+.2f}" for vp, r in zip([0.50,0.60,0.70], row)))

    # (3) short-borrow + long-only
    print("(3) short-borrow & long-only (net):")
    for ba in [0.0, 0.10, 0.30]:
        print(f"     L/S borrow {int(ba*100)}%/yr: {sharpe_t(T_(lowvol_pnl(base, borrow_ann=ba))):+.2f}")
    lo = lowvol_pnl(base, long_only=True)
    vlo = validate(lo, days, 30)
    print(f"     LONG-ONLY (no borrow): full {vlo['full']:+.2f} OOS {vlo['oos']:+.2f} CPCVmed {vlo['med']:+.2f} DSR {vlo['dsr']:.2f}")

    # (4) per-year (base L/S)
    print("(4) per-year (base L/S): " + " ".join(f"{y}:{sharpe_t(T_(base_pnl[year[1:]==y])):+.2f}" for y in range(2023,2027) if (year[1:]==y).sum()>40))

    # (5) correlation to crypto momentum book
    syms, cdays, cc, cqv, cf = pit_sweep.load()
    cw, _ = compute_weights(cc, cqv, cdays, CFG)
    cR = np.zeros_like(cc); cR[1:] = np.log(cc)[1:] - np.log(cc)[:-1]; cR = np.where(np.isfinite(cR), cR, 0.0)
    cff = np.where(np.isfinite(cf), cf, 0.0)
    cmom = np.sum(cw[:-1] * (cR - cff)[1:], axis=1)
    cmd = {int(cdays[1:][i]): cmom[i] for i in range(len(cmom))}
    emd = {int(days[1:][i]): base_pnl[i] for i in range(len(base_pnl))}
    common = sorted(set(cmd) & set(emd))
    a = np.array([cmd[d] for d in common]); b = np.array([emd[d] for d in common])
    m = np.isfinite(a) & np.isfinite(b)
    corr = float(np.corrcoef(a[m], b[m])[0, 1]) if m.sum() > 50 else float("nan")
    print(f"(5) corr to crypto-momentum book: {corr:+.2f}  (low => diversifier in a preferred market)")
    print("\nVERDICT: bootstrap frac>0~1 + grid broadly + long-only positive + DSR>0.5 + low crypto-corr = real 2nd sleeve.")


if __name__ == "__main__":
    main()