#!/usr/bin/env python3
"""Proper research: (a) robust momentum construction + (b) diversifying second edge.

All point-in-time (universe = daily top-50 by trailing dollar-volume, dead coins included) and
under REALISM stress (Reff winsorized ±20%/day, death-spiral last-5d excluded, 20bp cost) — we
optimize for robust/realistic performance, not headline.

(a) 6 pre-registered momentum constructions judged by STRESSED Sharpe + degradation ratio + every-year.
(b) canonical diversifiers (carry, TS-trend, long/short reversal, low-vol, size) judged by correlation
    to the momentum sleeve + COMBINED inverse-vol book Sharpe (a weak-but-uncorrelated sleeve can help).
"""
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pit_sweep as ps  # 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"
WINSOR, COST, TOPK = 0.20, 20.0, 50


def tr(lc, L, skip=0):
    out = np.full_like(lc, np.nan)
    if skip == 0:
        out[L:] = lc[L:] - lc[:-L]
    else:  # return from t-L-skip to t-skip (classic 12-1 skip)
        out[L + skip:] = lc[L:-skip] - lc[:-(L + skip)]
    return out


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 row_rank(x):  # per-row percentile rank in [-0.5,0.5], nan-aware
    out = np.full_like(x, np.nan)
    for t in range(x.shape[0]):
        m = np.isfinite(x[t])
        if m.sum() > 1:
            r = x[t, m].argsort().argsort().astype(float)
            out[t, m] = r / (m.sum() - 1) - 0.5
    return out


def main():
    syms, days, close, qv, fund = ps.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)
    year = (1970 + days / 365.25).astype(int)
    yrs = list(range(2020, 2027))

    # PIT universe + tradeable (death-spiral) mask
    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
    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_base = R - fund
    Reff_str = np.where(tradeable, np.clip(Reff_base, -WINSOR, WINSOR), 0.0)
    vol20 = roll(np.std, R, 20)

    def pnl(sig, reff, cost):
        s = sig.copy(); s[~univ] = np.nan
        return pnl_w(xs_weights(s), reff, cost_bp=cost)

    def peryear(p):
        return [sharpe_t(torch.tensor(p[year[1:] == y], device=DEV, dtype=torch.float64))
                if (year[1:] == y).sum() > 30 else float("nan") for y in yrs]

    # ---------- (a) momentum constructions ----------
    raw20 = tr(lc, 20)
    cons = {
        "raw_20": raw20,
        "rank_20": row_rank(raw20),
        "volscaled_20": raw20 / np.where(vol20 > 0, vol20, np.nan),
        "pathconsist_20": roll(np.mean, np.sign(R), 20),
        "skiprecent_20s3": tr(lc, 20, skip=3),
        "winsorinput_20": roll(np.sum, np.clip(R, -0.05, 0.05), 20),
    }
    print(f"\n===== (a) ROBUST MOMENTUM CONSTRUCTION (PIT top-{TOPK}, stress=winsor±{int(WINSOR*100)}%+death+{int(COST)}bp) =====")
    print(f"{'construction':>16} {'base_full':>9} {'str_full':>8} {'str_OOS':>7} {'str_CPCV':>8} {'ratio':>5} | per-year(stress) 2020..26")
    abest, abest_pnl, abest_score = None, None, -9
    for nm, sg in cons.items():
        pb = pnl(sg, Reff_base, 10.0); ps_ = pnl(sg, Reff_str, COST)
        sb = sharpe_t(torch.tensor(pb, device=DEV, dtype=torch.float64))
        v = validate(ps_, days, 6)
        ratio = v["full"] / sb if sb and not math.isnan(sb) and sb != 0 else float("nan")
        py = peryear(ps_)
        pys = " ".join(f"{p:>+5.2f}" if not math.isnan(p) else "  n/a" for p in py)
        nneg = sum(1 for p in py if not math.isnan(p) and p < 0)
        score = v["med"] - 0.3 * nneg  # robust: high stressed CPCV-med, few negative years
        print(f"{nm:>16} {sb:>+9.2f} {v['full']:>+8.2f} {v['oos']:>+7.2f} {v['med']:>+8.2f} {ratio:>5.2f} | {pys}")
        if score > abest_score:
            abest, abest_pnl, abest_score = nm, ps_, score
    print(f"  -> robust momentum sleeve = {abest}")

    # ---------- (b) diversifying second edge ----------
    sleeve = abest_pnl
    cand = {
        "carry_7": -roll(np.mean, fund, 7),
        "TStrend_30": np.sign(tr(lc, 30)) * np.where(univ, 1.0, np.nan),
        "LTreversal_180": -tr(lc, 180),
        "STreversal_2": -tr(lc, 2),
        "lowvol": -vol20,
        "size_small": -np.log(np.where(dv30 > 0, dv30, np.nan)),
    }
    print(f"\n===== (b) DIVERSIFYING EDGE vs momentum sleeve ({abest}) — stress applied =====")
    print(f"{'candidate':>16} {'standalone':>10} {'corr_to_mom':>11} {'combined':>9} | combined per-year 2020..26")
    sl = torch.tensor(sleeve, device=DEV, dtype=torch.float64)
    sl_sd = float(sl.std())
    combos = []
    for nm, sg in cand.items():
        pc = pnl(sg, Reff_str, COST)
        a, b = np.asarray(sleeve), np.asarray(pc)
        m = np.isfinite(a) & np.isfinite(b)
        corr = float(np.corrcoef(a[m], b[m])[0, 1]) if m.sum() > 50 else float("nan")
        sc = float(torch.tensor(pc, device=DEV, dtype=torch.float64).std())
        wm, wc = (1 / sl_sd), (1 / sc if sc > 0 else 0)
        comb = (wm * a + wc * b) / (wm + wc)
        vc = sharpe_t(torch.tensor(comb, device=DEV, dtype=torch.float64))
        py = peryear(comb)
        pys = " ".join(f"{p:>+5.2f}" if not math.isnan(p) else "  n/a" for p in py)
        sa = sharpe_t(torch.tensor(pc, device=DEV, dtype=torch.float64))
        print(f"{nm:>16} {sa:>+10.2f} {corr:>+11.2f} {vc:>+9.2f} | {pys}")
        combos.append((nm, corr, vc, pc))

    # ---------- final BOOK: momentum + diversifiers that lift combined Sharpe and are low-corr ----------
    base_sh = sharpe_t(sl)
    keep = [(nm, pc) for nm, corr, vc, pc in combos if vc > base_sh and (math.isnan(corr) or corr < 0.5)]
    print(f"\n===== BOOK = momentum + {[k for k,_ in keep]} (inverse-vol weighted, stress) =====")
    sleeves = [np.asarray(sleeve)] + [np.asarray(pc) for _, pc in keep]
    sds = [float(np.nanstd(s)) for s in sleeves]
    ws = np.array([1 / s if s > 0 else 0 for s in sds]); ws /= ws.sum()
    book = sum(w * s for w, s in zip(ws, sleeves))
    vb = validate(book, days, 6)
    py = peryear(book)
    pys = " ".join(f"{p:>+5.2f}" if not math.isnan(p) else "  n/a" for p in py)
    print(f"  momentum-only Sharpe (stress) = {base_sh:+.2f}")
    print(f"  BOOK: full {vb['full']:+.2f}  IS {vb['is_']:+.2f}  OOS {vb['oos']:+.2f}  CPCVmed {vb['med']:+.2f}  CPCV5% {vb['p5']:+.2f}")
    print(f"  BOOK per-year 2020..26: {pys}")
    print("\nVERDICT: book Sharpe > momentum-only with every-year positive => diversification works (deployable book).")


if __name__ == "__main__":
    main()