#!/usr/bin/env python3
"""Capstone research: (2) low-vol full validation battery + (1) per-coin slippage/capacity model.

Point-in-time universe (daily top-50 by trailing dollar-vol, dead coins in). Returns are
death-spiral-excluded (last 5 untradeable days zeroed) but NOT blanket-winsorized — the
slippage model replaces the crude cap with realistic, liquidity-aware cost:
   one-way cost_i = half_spread(ADV_i) + η·σ_i·sqrt(trade_notional_i / ADV_i)   (square-root impact)
Sweeping AUM gives the capacity curve (Sharpe vs size) for momentum, low-vol, and the book.
Low-vol gets the same battery momentum passed (window sweep, coin-bootstrap, per-year, IS/OOS).
"""
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 = 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 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)        # death-excl, NO blanket winsor

    def W(sig):
        s = sig.copy(); s[~univ] = np.nan
        return xs_weights(s)

    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]

    # ---------- (2) LOW-VOL BATTERY ----------
    print(f"\n===== (2) LOW-VOL VALIDATION BATTERY (PIT top-{TOPK}, death-excl, 10bp) =====")
    print(f"{'vol-window':>10} {'full':>6} {'IS':>6} {'OOS':>6} {'CPCVmed':>8} {'DSR':>5} | per-year 2020..26")
    for Wd in [10, 20, 30, 60]:
        sig = -np.nan_to_num(roll(np.std, R, Wd))
        pnl = pnl_w(W(sig), Reff, cost_bp=10)
        v = validate(pnl, days, 8)
        py = " ".join(f"{p:>+5.2f}" if not math.isnan(p) else "  n/a" for p in peryear(pnl))
        print(f"{Wd:>10} {v['full']:>+6.2f} {v['is_']:>+6.2f} {v['oos']:>+6.2f} {v['med']:>+8.2f} {v['dsr']:>5.2f} | {py}")
    # coin-bootstrap on lowvol(20)
    rng = np.random.default_rng(7)
    lvsig = -vol20
    sh = []
    for _ in range(200):
        cols = rng.choice(N, size=max(N // 2, 10), replace=False)
        sub = np.full((T, N), np.nan); sub[:, cols] = lvsig[:, cols]
        s = sharpe_t(torch.tensor(pnl_w(W(sub), Reff, cost_bp=10), device=DEV, dtype=torch.float64))
        if not math.isnan(s):
            sh.append(s)
    sh = np.array(sh)
    print(f"  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}")

    # ---------- (1) SLIPPAGE / CAPACITY MODEL ----------
    w_mom = W(tr(lc, 20))           # raw 20d momentum (diversifies w/ lowvol, corr -0.13)
    w_lv = W(-vol20)
    w_book = 0.5 * (w_mom + w_lv)   # netting preserved (offsetting positions reduce turnover)

    def net_sharpe(w, AUM, eta=1.0):
        turn = np.abs(w[1:] - w[:-1])                       # fraction of AUM traded per coin
        adv = dv30[1:]; sg = vol20[1:]
        hs = np.clip(30.0 / np.sqrt(np.maximum(adv, 1.0) / 1e6), 1.0, 30.0) / 1e4   # half-spread frac
        part = np.where(adv > 0, turn * AUM / adv, 0.0)
        impact = eta * sg * np.sqrt(np.clip(part, 0, None))
        cost = np.sum(turn * (hs + impact), axis=1)
        gross = np.sum(w[:-1] * Reff[1:], axis=1)
        return sharpe_t(torch.tensor(gross - cost, device=DEV, dtype=torch.float64))

    print(f"\n===== (1) SLIPPAGE / CAPACITY — Sharpe vs AUM (sqrt-impact η=1, liquidity-scaled spread) =====")
    aums = [1e6, 5e6, 2e7, 5e7, 2e8, 5e8]
    print(f"{'sleeve':>12} " + " ".join(f"{'$'+(f'{a/1e6:.0f}M' if a<1e9 else f'{a/1e9:.1f}B'):>8}" for a in aums))
    for nm, w in [("momentum", w_mom), ("lowvol", w_lv), ("BOOK", w_book)]:
        row = [net_sharpe(w, a) for a in aums]
        print(f"{nm:>12} " + " ".join(f"{r:>+8.2f}" for r in row))
    # frictionless reference (death-excl only, no cost)
    print(f"{'(grossSR)':>12} " + " ".join(
        f"{sharpe_t(torch.tensor(np.sum(w_book[:-1]*Reff[1:],axis=1),device=DEV,dtype=torch.float64)):>+8.2f}" for _ in [0]) +
        "  <- book gross (no cost)")
    print("\nVERDICT: Sharpe at $5-50M AUM = realistic deployable number; where it decays = capacity ceiling.")


if __name__ == "__main__":
    main()