#!/usr/bin/env python3
"""Diversifier hunt: directional TIME-SERIES trend (CTA-style) vs market-neutral XS momentum.

XS momentum (the validated edge) is market-NEUTRAL relative-value. A proper TS-trend sleeve is
DIRECTIONAL — net-long when coins trend up, net-SHORT when they trend down — so it should profit
in sustained bears (2022) when XS momentum is weak: the classic crisis-alpha complement.
Tests TS-trend standalone battery (lookback, per-year, bootstrap), its CORRELATION to the
momentum sleeve, and the COMBINED book. PIT universe, death-excl, realistic.
"""
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"
TOPK = 30


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 tr(lc, L):
    out = np.full_like(lc, np.nan); out[L:] = lc[L:] - lc[:-L]; 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)
    vol20 = np.nan_to_num(roll(np.std, R, 20))
    year = (1970 + days / 365.25).astype(int); yrs = list(range(2020, 2027))
    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 = np.where(tradeable, R - fund, 0.0)
    invvol = np.where(vol20 > 0, 1.0 / vol20, 0.0)

    def wn(sig):  # market-neutral (demeaned) unit-gross
        s = sig.copy(); s[~univ] = np.nan
        return xs_weights(s)

    def wdir(sig):  # DIRECTIONAL unit-gross (NOT demeaned) — net long/short exposure
        s = np.where(univ, sig, 0.0)
        g = np.sum(np.abs(s), axis=1, keepdims=True); g[g == 0] = 1
        return s / g

    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]

    mom_pnl = pnl_w(wn(tr(lc, 20)), Reff, cost_bp=10)
    print(f"\n===== DIVERSIFIER HUNT: directional TS-trend vs XS momentum (PIT top-{TOPK}) =====")
    print("XS momentum sleeve per-year:  " + " ".join(f"{p:>+5.2f}" for p in peryear(mom_pnl)) +
          f"   (full {sharpe_t(torch.tensor(mom_pnl,device=DEV,dtype=torch.float64)):+.2f})")
    print(f"\n[TS-trend] DIRECTIONAL sign(trail_L)*invvol, net exposure")
    print(f"{'lookback':>9} {'full':>6} {'IS':>6} {'OOS':>6} {'CPCVmed':>8} {'corr_mom':>8} {'combined':>9} | per-year 2020..26 (TS)")
    for L in [20, 30, 60, 90]:
        ts = wdir(np.sign(tr(lc, L)) * invvol)
        ts_pnl = pnl_w(ts, Reff, cost_bp=10)
        v = validate(ts_pnl, days, 8)
        a, b = mom_pnl, ts_pnl
        m = np.isfinite(a) & np.isfinite(b)
        corr = float(np.corrcoef(a[m], b[m])[0, 1])
        sm = float(np.nanstd(a)); st = float(np.nanstd(b))
        comb = ((1 / sm) * a + (1 / st) * b) / (1 / sm + 1 / st)
        vc = sharpe_t(torch.tensor(comb, device=DEV, dtype=torch.float64))
        py = " ".join(f"{p:>+5.2f}" if not math.isnan(p) else "  n/a" for p in peryear(ts_pnl))
        print(f"{L:>9} {v['full']:>+6.2f} {v['is_']:>+6.2f} {v['oos']:>+6.2f} {v['med']:>+8.2f} {corr:>+8.2f} {vc:>+9.2f} | {py}")

    # bootstrap on TS-trend(30)
    rng = np.random.default_rng(11)
    base = np.sign(tr(lc, 30)) * invvol
    sh = []
    for _ in range(200):
        cols = rng.choice(N, size=max(N // 2, 10), replace=False)
        sub = np.zeros((T, N)); sub[:, cols] = base[:, cols]
        s = sharpe_t(torch.tensor(pnl_w(wdir(np.where(univ, sub, 0)), Reff, cost_bp=10), device=DEV, dtype=torch.float64))
        if not math.isnan(s):
            sh.append(s)
    sh = np.array(sh)
    print(f"\nTS-trend(30) coin-bootstrap(200): mean {sh.mean():+.2f} median {np.median(sh):+.2f} 5th {np.percentile(sh,5):+.2f} frac>0 {np.mean(sh>0):.2f}")

    # best combined book per-year
    ts30 = pnl_w(wdir(np.sign(tr(lc, 30)) * invvol), Reff, cost_bp=10)
    sm, st = float(np.nanstd(mom_pnl)), float(np.nanstd(ts30))
    book = ((1 / sm) * mom_pnl + (1 / st) * ts30) / (1 / sm + 1 / st)
    vb = validate(book, days, 8)
    print(f"\nBOOK = XS-momentum + TS-trend(30) (inv-vol): full {vb['full']:+.2f} IS {vb['is_']:+.2f} OOS {vb['oos']:+.2f} CPCVmed {vb['med']:+.2f} 5% {vb['p5']:+.2f}")
    print("BOOK per-year 2020..26: " + " ".join(f"{p:>+5.2f}" for p in peryear(book)))
    print("\nVERDICT: TS-trend real (battery) + low corr + 2022 positive + lifts book => genuine diversifier.")


if __name__ == "__main__":
    main()