#!/usr/bin/env python3
"""The decisive gate both the critic and the exhaustion-researcher named: does the
exhaustion-short carry MARGINAL ALPHA over plain residual momentum, or is it collinear?

Build, on the PIT crypto panel: (1) residual-momentum L/S book [validated], (2) exhaustion
L/S book (z-sum of: neg residual-mom, fading dollar-volume rank, extreme funding, drawdown).
Then: correlation, and regress exhaustion daily returns on momentum daily returns -> the
intercept (alpha) is what matters. High corr + ~0 alpha => momentum re-spelled => drop it.
Also A/B the squeeze veto (no short when funding<=0). Realistic: Reff = return - funding
(shorts pay positive funding), death-excl, 10bp.
"""
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pit_sweep  # noqa: E402
from signal_sweep import xs_weights, pnl_w, validate, sharpe_t  # noqa: E402
import torch  # noqa: E402

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


def roll(fn, X, L):
    out = np.full_like(X, np.nan)
    for t in range(L, len(X)):
        out[t] = fn(X[t - L:t], axis=0)
    return out


def zc(x):
    mu = np.nanmean(x, axis=1, keepdims=True); sd = np.nanstd(x, axis=1, keepdims=True)
    return np.nan_to_num((x - mu) / np.where(sd > 0, sd, 1))


def main():
    syms, days, close, qv, fund = pit_sweep.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)
    fund = np.where(np.isfinite(fund), fund, 0.0)
    qv = np.nan_to_num(qv)
    dv30 = roll(np.mean, qv, 30)
    tradeable = np.ones((T, N), bool)
    for j in range(N):
        idx = np.where(np.isfinite(close[:, j]))[0]
        if len(idx):
            tradeable[max(0, idx[-1] - 4):idx[-1] + 1, j] = False
    Reff = np.where(tradeable, R - fund, 0.0)
    univ = np.zeros((T, N), bool)
    for t in range(T):
        elig = np.where((dv30[t] > 0) & np.isfinite(close[t]))[0]
        if len(elig):
            univ[t, elig[np.argsort(-dv30[t, elig])[:TOPK]]] = True
    year = (1970 + days / 365.25).astype(int)

    # residual (beta-stripped) momentum
    mkt = np.array([R[t][univ[t]].mean() if univ[t].any() else 0.0 for t in range(T)])
    beta = np.zeros((T, N)); Wb = 60
    for t in range(Wb, T):
        mw = mkt[t - Wb:t]; vb = mw.var() + 1e-12
        beta[t] = ((R[t - Wb:t] * mw[:, None]).mean(0) - R[t - Wb:t].mean(0) * mw.mean()) / vb
    rcum = np.cumsum(R - beta * mkt[:, None], axis=0)
    resid_mom = np.full((T, N), np.nan); resid_mom[20:] = rcum[20:] - rcum[:-20]

    # exhaustion terms (all oriented so HIGH => short candidate)
    dvr = np.full((T, N), np.nan)                                # cross-sectional dollar-vol rank (0..1)
    for t in range(T):
        e = np.where(univ[t])[0]
        if len(e) > 1:
            r = dv30[t, e].argsort().argsort().astype(float) / (len(e) - 1)
            dvr[t, e] = r
    dvr_slope = np.full((T, N), np.nan); dvr_slope[30:] = dvr[30:] - dvr[:-30]   # falling rank = fading
    fund7 = roll(np.mean, fund, 7)                               # crowded-long carry
    rmax = roll(np.max, lc, 90); dd = lc - rmax                  # drawdown-from-ATH (<=0)

    exh = zc(-resid_mom) + zc(-dvr_slope) + zc(fund7) + zc(-dd)  # additive, equal-weight (no fitting)

    def book(sig, veto_short=None):
        s = sig.copy(); s[~univ] = np.nan
        w = xs_weights(s)
        if veto_short is not None:
            w = np.where((w < 0) & veto_short, 0.0, w)           # drop shorts where veto true
        return pnl_w(w, Reff, cost_bp=10)

    pnl_mom = book(resid_mom)
    pnl_exh = book(-exh)                                         # long low-exhaustion / short high
    squeeze_veto = fund <= 0                                     # don't short crowded-short (squeeze fuel)
    pnl_exh_veto = book(-exh, veto_short=squeeze_veto)

    def stats(p):
        v = validate(p, days, 50); return v
    T_ = lambda x: torch.tensor(x[np.isfinite(x)], device=DEV, dtype=torch.float64)

    print(f"\n===== EXHAUSTION-SHORT MARGINAL-ALPHA GATE (PIT top{TOPK}, death-excl, 10bp, deflate N=50) =====")
    for nm, p in [("residual_momentum", pnl_mom), ("exhaustion", pnl_exh), ("exhaustion+veto", pnl_exh_veto)]:
        v = stats(p)
        print(f"  {nm:>18}: full {v['full']:+.2f}  IS {v['is_']:+.2f}  OOS {v['oos']:+.2f}  CPCVmed {v['med']:+.2f}  DSR {v['dsr']:.2f}")

    # THE decisive test: regress exhaustion returns on momentum returns -> marginal alpha
    a, b = pnl_mom, pnl_exh
    m = np.isfinite(a) & np.isfinite(b)
    x, y = a[m], b[m]
    corr = float(np.corrcoef(x, y)[0, 1])
    beta1 = float(np.cov(x, y)[0, 1] / (np.var(x) + 1e-12))
    resid = y - beta1 * x
    alpha_daily = float(resid.mean())
    alpha_ann_sr = alpha_daily / (resid.std() + 1e-12) * math.sqrt(365)
    alpha_t = alpha_daily / (resid.std() / math.sqrt(len(resid)) + 1e-12)
    print(f"\n  REGRESS exhaustion ~ momentum:  corr {corr:+.2f}  beta {beta1:+.2f}")
    print(f"  marginal alpha: ann-Sharpe {alpha_ann_sr:+.2f}  t-stat {alpha_t:+.2f}  (>2 = real distinct edge)")

    # does combining beat momentum alone (OOS)?
    sm, se = np.nanstd(a), np.nanstd(b)
    comb = (np.nan_to_num(a) / sm + np.nan_to_num(b) / se)
    vc = stats(comb); vm = stats(pnl_mom)
    print(f"\n  combined(mom+exh) OOS {vc['oos']:+.2f} vs momentum-alone OOS {vm['oos']:+.2f}  "
          f"=> {'ADDS' if vc['oos'] > vm['oos'] + 0.1 else 'no improvement'}")
    print("\nVERDICT: high corr + alpha t<2 + no OOS improvement => exhaustion is momentum re-spelled -> DROP (ship Comp-1).")


if __name__ == "__main__":
    main()