#!/usr/bin/env python3
"""Surfer PoC — crypto cross-sectional momentum (the validated edge), end-to-end.

Modes:
  backtest  — run the strategy over the cached point-in-time history; print honest stats
              (gross/net Sharpe at AUM, OOS, CPCV-median, DSR, per-year, max-DD, turnover, capacity curve)
  regime    — diagnostic: does any regime feature predict when momentum works? (overlay is OFF by default)
  paper     — forward paper-trade: fetch live data, compute today's target weights, log intended trades,
              mark the paper book (price + funding), persist state. The only true out-of-sample test.
  status    — print the paper book + equity curve summary.

Design: ONE signal function (`compute_weights`) drives BOTH backtest and live (no backtest/live skew).
No lookahead (signal uses closed bars <= t; weights apply to t+1; rebalance keyed on calendar day%K).
Funding is P&L (longs pay, shorts earn). Realistic sqrt-impact cost. Validated edge only; regime overlay
stays a diagnostic until separately confirmed.
"""
import argparse
import json
import math
import os
import sys
import time
import urllib.request

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from signal_sweep import xs_weights, validate, sharpe_t  # noqa: E402
import pit_sweep  # noqa: E402  (load() of the cached PIT universe)
import torch  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"
DAY_MS = 86_400_000
_REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))  # cwd-independent (cron-safe)
STATE = os.path.join(_REPO, "data/surfer/poc_state.json")

CFG = dict(
    lookback=20,        # XS momentum horizon (days)
    beta_win=60,        # rolling window for market-beta (residual momentum strips this)
    topk=30,            # universe = top-K liquid by trailing dollar-vol (breadth/liquidity optimum)
    rebal_k=7,          # rebalance every 7 calendar days (weekly)
    smooth_span=5,      # EWMA weight-smoothing span (cuts turnover + whipsaw)
    vol_win=30,         # trailing vol window (for sizing / regime)
    cost_aum=2e7,       # AUM the slippage model is priced at ($20M)
    eta=1.0,            # sqrt market-impact coefficient (conservative)
    dd_breaker=-0.20,   # cut gross exposure if rolling drawdown worse than this
    vrp_risk_frac=0.30, # VRP diversifier sleeve risk budget (small — short-vol tail)
    vrp_rise_k=5,       # don't sell vol when DVOL rose over last K days (tail gate)
    vrp_level_win=60,   # sell full size only when DVOL < trailing-win median
)

# ---------------------------------------------------------------- helpers
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 trailing(lc, L):
    out = np.full_like(lc, np.nan); out[L:] = lc[L:] - lc[:-L]; return out


def ewma_rows(W, span):
    if span <= 1:
        return W
    a = 2.0 / (span + 1)
    out = W.copy()
    for t in range(1, len(W)):
        out[t] = a * W[t] + (1 - a) * out[t - 1]
    return out


# ---------------------------------------------------------------- SHARED SIGNAL (backtest == live)
def compute_weights(close, qvol, days, cfg):
    """Panel -> held target weights [T,N] (market-neutral, weekly, smoothed). No lookahead.
    Live takes the last row; backtest uses all rows. Identical code path."""
    T, N = close.shape
    lc = np.log(close)
    dv = roll(np.mean, np.nan_to_num(qvol), 30)
    univ = np.zeros((T, N), bool)
    for t in range(T):
        elig = np.where((dv[t] > 0) & np.isfinite(close[t]))[0]
        if len(elig):
            univ[t, elig[np.argsort(-dv[t, elig])[:cfg["topk"]]]] = True
    # residual (beta-stripped) momentum: strip each coin's beta to the equal-weight market,
    # rank the residual trend (Phase-1: +0.72 vs +0.57 raw — cleaner base edge). No lookahead.
    R = np.zeros((T, N)); R[1:] = lc[1:] - lc[:-1]; R = np.where(np.isfinite(R), R, 0.0)
    mkt = np.array([R[t][univ[t]].mean() if univ[t].any() else 0.0 for t in range(T)])
    Wb = cfg["beta_win"]; beta = np.zeros((T, N))
    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)
    L = cfg["lookback"]
    sig = np.full((T, N), np.nan); sig[L:] = rcum[L:] - rcum[:-L]
    sig[~univ] = np.nan
    wt = xs_weights(sig)                              # daily target (market-neutral, unit gross)
    wt = ewma_rows(wt, cfg["smooth_span"])           # smooth
    held = wt.copy()                                 # weekly rebalance on FIXED calendar phase (day%K==0)
    K = cfg["rebal_k"]                               # epoch day0=Thu → rebalances every Thursday, live-consistent
    last = 0
    for t in range(T):
        if days[t] % K == 0:
            last = t
        held[t] = wt[last]
    return held, univ


def slippage_net(w, Reff, dv, vol, days, cfg):
    """Net daily PnL series under sqrt market-impact + liquidity-scaled spread at cfg['cost_aum']."""
    AUM = cfg["cost_aum"]
    turn = np.abs(w[1:] - w[:-1])
    adv = np.nan_to_num(dv[1:]); sg = np.nan_to_num(vol[1:])
    hs = np.clip(30.0 / np.sqrt(np.maximum(adv, 1.0) / 1e6), 1.0, 30.0) / 1e4
    part = np.where(adv > 0, turn * AUM / adv, 0.0)
    cost = np.sum(turn * (hs + cfg["eta"] * sg * np.sqrt(np.clip(part, 0, None))), axis=1)
    return np.sum(w[:-1] * Reff[1:], axis=1) - cost


def max_drawdown(pnl):
    eq = np.cumsum(np.nan_to_num(pnl)); peak = np.maximum.accumulate(eq)
    return float((eq - peak).min())


# ---------------------------------------------------------------- VRP SLEEVE (tail-managed short-vol diversifier)
def _dvol(ccy, refresh):
    cache = os.path.join(_REPO, "data/surfer/dvol", f"{ccy}.json")
    d = {int(k): v for k, v in json.load(open(cache)).items()} if os.path.exists(cache) else {}
    if refresh:
        end = int(time.time() * 1000); start = end - 200 * DAY_MS
        try:
            rows = _get(f"https://www.deribit.com/api/v2/public/get_volatility_index_data?currency={ccy}"
                        f"&start_timestamp={start}&end_timestamp={end}&resolution=86400").get("result", {}).get("data", [])
            for r in rows:
                d[int(r[0]) // DAY_MS] = float(r[4])
            os.makedirs(os.path.dirname(cache), exist_ok=True)
            json.dump({str(k): v for k, v in d.items()}, open(cache, "w"))
        except Exception:
            pass
    return d


def vrp_ccy(iv, close, cfg):
    """Tail-managed short-variance daily P&L per ccy. Gate uses DVOL up to t-1 (no lookahead)."""
    days = sorted(set(iv) & set(close)); pnl = {}
    w0 = max(cfg["vrp_level_win"], cfg["vrp_rise_k"]) + 1
    for i in range(w0, len(days)):
        d0, d1 = days[i - 1], days[i]
        if close[d0] <= 0 or close[d1] <= 0:
            continue
        r = math.log(close[d1] / close[d0]); ivl = iv[d0] / 100.0
        raw = ivl * ivl / 365.0 - r * r                          # collect implied var, pay realized
        rising = iv[d0] > iv[days[i - 1 - cfg["vrp_rise_k"]]]     # don't sell into rising vol
        win = [iv[days[j]] for j in range(i - 1 - cfg["vrp_level_win"], i - 1)]
        level = 1.0 if iv[d0] < float(np.median(win)) else 0.5    # full in calm, half elevated
        pnl[d1] = (0.0 if rising else 1.0) * level * raw
    return pnl


def vrp_book_series(cfg, btc_close, eth_close, refresh):
    b = vrp_ccy(_dvol("BTC", refresh), btc_close, cfg)
    e = vrp_ccy(_dvol("ETH", refresh), eth_close, cfg)
    days = sorted(set(b) | set(e))
    return {d: float(np.nanmean([b.get(d, np.nan), e.get(d, np.nan)])) for d in days}


def _closes_from(syms, days, close, sym):
    if sym not in syms:
        return {}
    j = syms.index(sym)
    return {int(days[t]): float(close[t, j]) for t in range(len(days)) if np.isfinite(close[t, j])}


def _vrp_calib(cfg):
    """From cached history: (momentum daily vol, VRP daily vol) — to size VRP to its risk budget."""
    syms, days, close, qv, fund = pit_sweep.load()
    R = np.zeros_like(close); R[1:] = np.log(close)[1:] - np.log(close)[:-1]; R = np.where(np.isfinite(R), R, 0.0)
    cf = np.where(np.isfinite(fund), fund, 0.0)
    w, _ = compute_weights(close, qv, days, cfg)
    mom = np.sum(w[:-1] * (R - cf)[1:], axis=1)
    vrpb = vrp_book_series(cfg, _closes_from(syms, days, close, "BTCUSDT"),
                           _closes_from(syms, days, close, "ETHUSDT"), refresh=False)
    return float(np.nanstd(mom)), float(np.nanstd(np.array(list(vrpb.values()))))


# ---------------------------------------------------------------- BACKTEST
def backtest(cfg):
    syms, days, close, qv, fund = pit_sweep.load()
    T, N = close.shape
    R = np.zeros((T, N)); R[1:] = np.log(close)[1:] - np.log(close)[:-1]
    R = np.where(np.isfinite(R), R, 0.0)
    fund = np.where(np.isfinite(fund), fund, 0.0)
    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)
    dv = roll(np.mean, np.nan_to_num(qv), 30)
    vol = roll(np.std, R, cfg["vol_win"])
    w, univ = compute_weights(close, qv, days, cfg)
    gross = np.sum(w[:-1] * Reff[1:], axis=1)
    net = slippage_net(w, Reff, dv, vol, days, cfg)
    year = (1970 + days / 365.25).astype(int)
    v = validate(net, days, 30)
    gsr = sharpe_t(torch.tensor(gross, device=DEV, dtype=torch.float64))
    turn = float(np.mean(np.sum(np.abs(w[1:] - w[:-1]), axis=1))) * 100
    print(f"\n===== SURFER PoC BACKTEST — XS momentum (top{cfg['topk']}, {cfg['lookback']}d, weekly, smooth{cfg['smooth_span']}) =====")
    print(f"coins(ever)={N}  days={T}  universe/day~{int(univ.sum(1).mean())}  cost-AUM=${cfg['cost_aum']/1e6:.0f}M  turnover/day={turn:.1f}%")
    print(f"gross Sharpe {gsr:+.2f} | NET Sharpe {v['full']:+.2f}  IS {v['is_']:+.2f}  OOS {v['oos']:+.2f}  CPCVmed {v['med']:+.2f}  DSR {v['dsr']:.2f}")
    print(f"max-DD (sum-rets) {max_drawdown(net):+.2f}")
    py = [sharpe_t(torch.tensor(net[year[1:] == y], device=DEV, dtype=torch.float64)) if (year[1:] == y).sum() > 30 else float('nan') for y in range(2020, 2027)]
    print("per-year 2020..26: " + " ".join(f"{p:+.2f}" if not math.isnan(p) else " n/a" for p in py))
    print("capacity:  " + "  ".join(
        f"${a/1e6:.0f}M={sharpe_t(torch.tensor(slippage_net(w,Reff,dv,vol,days,{**cfg,'cost_aum':a}),device=DEV,dtype=torch.float64)):+.2f}"
        for a in [1e6, 5e6, 2e7, 5e7, 2e8]))
    # ---- VRP diversifier sleeve + combined two-sleeve book ----
    vrpb = vrp_book_series(cfg, _closes_from(syms, days, close, "BTCUSDT"),
                           _closes_from(syms, days, close, "ETHUSDT"), refresh=False)
    mom_by = {int(days[1:][i]): net[i] for i in range(len(net))}
    common = sorted(set(vrpb) & set(mom_by))
    if common:
        mo = np.array([mom_by[d] for d in common]); vr = np.array([vrpb[d] for d in common])
        mu, vu, f = np.nanstd(mo), np.nanstd(vr), cfg["vrp_risk_frac"]
        comb = (mo / mu) * (1 - f) + (vr / vu) * f
        cy = (1970 + np.array(common) / 365.25).astype(int)
        T = lambda x: torch.tensor(x, device=DEV, dtype=torch.float64)
        print(f"VRP sleeve (tail-managed) Sharpe {sharpe_t(T(vr)):+.2f}  corr-to-momentum {np.corrcoef(mo,vr)[0,1]:+.2f}")
        print(f"COMBINED BOOK (momentum {int((1-f)*100)}% + VRP {int(f*100)}% risk) Sharpe {sharpe_t(T(comb)):+.2f}")
        print("combined per-year: " + " ".join(
            f"{y}:{sharpe_t(T(comb[cy==y])):+.2f}" for y in range(2021, 2027) if (cy == y).sum() > 30))


# ---------------------------------------------------------------- REGIME DIAGNOSTIC (overlay default OFF)
def regime(cfg):
    syms, days, close, qv, fund = pit_sweep.load()
    T, N = close.shape
    R = np.zeros((T, N)); R[1:] = np.log(close)[1:] - np.log(close)[:-1]; R = np.where(np.isfinite(R), R, 0.0)
    fund = np.where(np.isfinite(fund), fund, 0.0)
    w, univ = compute_weights(close, qv, days, cfg)
    book = np.sum(w[:-1] * (np.where(np.isfinite(R), R, 0.0) - fund)[1:], axis=1)   # daily book return
    lc = np.log(close)
    disp = np.array([np.nanstd(trailing(lc, cfg["lookback"])[t][univ[t]]) if univ[t].any() else np.nan for t in range(T)])
    mvol = np.nanmedian(roll(np.std, R, cfg["vol_win"]), axis=1)
    mtrend = np.array([np.nanmedian(np.abs(trailing(lc, cfg["lookback"])[t][univ[t]])) if univ[t].any() else np.nan for t in range(T)])
    feats = {"xs_dispersion": disp[:-1], "median_vol": mvol[:-1], "abs_trend": mtrend[:-1]}
    split = int(0.7 * len(book))
    print("\n===== REGIME DIAGNOSTIC — does a feature predict next-day book return? (overlay OFF until confirmed) =====")
    print(f"{'feature':>16} {'IC_all':>7} {'IC_IS':>7} {'IC_OOS':>7}")
    for nm, f in feats.items():
        f = f[1:]; b = book[1:]                     # feature_t -> book return_{t+1}
        m = np.isfinite(f) & np.isfinite(b)
        def ic(mask):
            mm = m & mask
            return float(np.corrcoef(f[mm], b[mm])[0, 1]) if mm.sum() > 50 else float("nan")
        idx = np.arange(len(f))
        print(f"{nm:>16} {ic(np.ones_like(m)):>+7.3f} {ic(idx < split):>+7.3f} {ic(idx >= split):>+7.3f}")
    print("Significant + IS/OOS-consistent IC => regime sizing worth building. Otherwise momentum runs flat-sized.")


# ---------------------------------------------------------------- LIVE (paper)
def _get(u):
    return json.load(urllib.request.urlopen(urllib.request.Request(u, headers={"User-Agent": "curl/8"}), timeout=30))


def _live_panel(cfg):
    info = _get("https://fapi.binance.com/fapi/v1/exchangeInfo")
    perps = {s["symbol"] for s in info["symbols"] if s.get("contractType") == "PERPETUAL"
             and s.get("quoteAsset") == "USDT" and s.get("status") == "TRADING"}
    tick = _get("https://fapi.binance.com/fapi/v1/ticker/24hr")
    vol = {t["symbol"]: float(t["quoteVolume"]) for t in tick if t["symbol"] in perps and t["symbol"].isascii()}
    syms = sorted(vol, key=lambda s: -vol[s])[:max(cfg["topk"] * 2, 60)]   # wide net; signal picks top-K
    need = max(cfg["lookback"] + cfg["vol_win"], cfg["vrp_level_win"] + cfg["vrp_rise_k"],
               cfg["beta_win"] + cfg["lookback"]) + 12  # cover VRP gate + residual-momentum beta window
    closes, qvols, funds, keep = {}, {}, {}, []
    for s in syms:
        k = _get(f"https://fapi.binance.com/fapi/v1/klines?symbol={s}&interval=1d&limit={need}")
        if len(k) < need - 2:
            continue
        d = np.array([int(r[0]) // DAY_MS for r in k])
        closes[s] = (d, np.array([float(r[4]) for r in k]), np.array([float(r[7]) for r in k]))
        fr = _get(f"https://fapi.binance.com/fapi/v1/fundingRate?symbol={s}&limit=10")
        funds[s] = float(fr[-1]["fundingRate"]) * 3 if fr else 0.0   # ~daily funding (3x 8h)
        keep.append(s)
        time.sleep(0.05)
    days = np.array(sorted(set().union(*[set(closes[s][0].tolist()) for s in keep])))
    di = {int(v): i for i, v in enumerate(days.tolist())}
    T, N = len(days), len(keep)
    close = np.full((T, N), np.nan); qv = np.full((T, N), np.nan)
    for j, s in enumerate(keep):
        d, c, q = closes[s]
        for i in range(len(d)):
            close[di[int(d[i])], j] = c[i]; qv[di[int(d[i])], j] = q[i]
    return keep, days, close, qv, np.array([funds[s] for s in keep])


def _load_state():
    if os.path.exists(STATE):
        return json.load(open(STATE))
    return dict(inception=None, last_mark_day=None, positions={}, prices={}, equity=1.0, log=[])


def _save_state(st):
    tmp = STATE + ".tmp"
    json.dump(st, open(tmp, "w"), indent=1)
    os.replace(tmp, STATE)


def paper(cfg):
    keep, days, close, qv, fund_now = _live_panel(cfg)
    w, univ = compute_weights(close, qv, days, cfg)
    target = {keep[j]: float(w[-1, j]) for j in range(len(keep)) if abs(w[-1, j]) > 1e-6}
    last_close = {keep[j]: float(close[-1, j]) for j in range(len(keep)) if np.isfinite(close[-1, j])}
    today = int(days[-1])
    st = _load_state()
    if st["inception"] is None:
        st["inception"] = today
    st.setdefault("vrp_equity", 1.0); st.setdefault("combined_equity", 1.0)
    if "vrp_vu" not in st:
        try:
            st["vrp_mu"], st["vrp_vu"] = _vrp_calib(cfg)
        except Exception:
            st["vrp_mu"], st["vrp_vu"] = 0.0, 1.0
    vrpb = vrp_book_series(cfg, _closes_from(keep, days, close, "BTCUSDT"),
                           _closes_from(keep, days, close, "ETHUSDT"), refresh=True)
    vrp_today = vrpb.get(today, None)
    # 1) MARK both sleeves to latest prices (+ funding) since last mark
    if st["positions"] and st["last_mark_day"] != today:
        ret = 0.0
        for s, pos in st["positions"].items():
            p0 = st["prices"].get(s); p1 = last_close.get(s)
            if p0 and p1 and p0 > 0:
                ret += pos * (p1 / p0 - 1.0)
            ret -= pos * fund_now[keep.index(s)] if s in keep else 0.0   # long pays funding
        f = cfg["vrp_risk_frac"]
        scaled = (vrp_today / st["vrp_vu"]) * st["vrp_mu"] if (vrp_today is not None and st["vrp_vu"] > 0) else 0.0
        comb_ret = (1 - f) * ret + f * scaled                            # VRP scaled to momentum vol, f risk
        st["equity"] *= (1.0 + ret)
        st["vrp_equity"] *= (1.0 + scaled)
        st["combined_equity"] *= (1.0 + comb_ret)
        st["log"].append({"day": today, "mom_ret": round(ret, 6), "vrp_ret": round(scaled, 6),
                          "comb_ret": round(comb_ret, 6), "combined": round(st["combined_equity"], 5)})
    # 2) REBALANCE on the weekly boundary
    rebal = (st["last_mark_day"] is None) or (today % cfg["rebal_k"] == 0)   # fixed weekly phase (Thursdays)
    trades = []
    if rebal:
        cur = st["positions"]
        allk = set(cur) | set(target)
        for s in sorted(allk):
            d = target.get(s, 0.0) - cur.get(s, 0.0)
            if abs(d) > 1e-4:
                trades.append({"sym": s, "side": "BUY" if d > 0 else "SELL", "dweight": round(d, 4)})
        st["positions"] = target
        st["prices"] = last_close
    st["last_mark_day"] = today
    _save_state(st)
    print(f"\n===== SURFER PoC PAPER — day {today} (universe {len(keep)} perps) =====")
    print(f"momentum {st['equity']:.4f} | VRP {st['vrp_equity']:.4f} | COMBINED BOOK {st['combined_equity']:.4f}  "
          f"(inception day {st['inception']}, {len(st['log'])} marks)")
    print(f"VRP today: {'flat (gated)' if (vrp_today is not None and abs(vrp_today)<1e-9) else ('%.5f'%vrp_today if vrp_today is not None else 'n/a')}   "
          f"rebalance: {rebal}   intended trades: {len(trades)}")
    longs = sorted([(s, x) for s, x in target.items() if x > 0], key=lambda z: -z[1])[:8]
    shorts = sorted([(s, x) for s, x in target.items() if x < 0], key=lambda z: z[1])[:8]
    print("top longs : " + ", ".join(f"{s} {x:+.3f}" for s, x in longs))
    print("top shorts: " + ", ".join(f"{s} {x:+.3f}" for s, x in shorts))
    for t in trades[:12]:
        print(f"   {t['side']:>4} {t['sym']:>12} dW {t['dweight']:+.4f}")
    print("(paper only — no real orders. Re-run weekly to forward-test edge persistence.)")


def status(cfg):
    st = _load_state()
    print(f"\n===== SURFER PoC STATUS =====")
    if st["inception"] is None:
        print("no paper state yet — run: python scripts/surfer/surfer_poc.py paper"); return
    print(f"inception day {st['inception']}  last mark {st['last_mark_day']}  marks {len(st['log'])}")
    print(f"equity: momentum {st['equity']:.4f} | VRP {st.get('vrp_equity',1.0):.4f} | COMBINED {st.get('combined_equity',1.0):.4f}")
    if st["log"]:
        def srp(key):
            r = np.array([m[key] for m in st["log"] if key in m])
            return r.mean() / (r.std() + 1e-9) * math.sqrt(252) if len(r) > 5 else float("nan")
        print(f"annual-ish Sharpe (live): momentum {srp('mom_ret'):+.2f} | VRP {srp('vrp_ret'):+.2f} | COMBINED {srp('comb_ret'):+.2f}")
    print(f"current book: {len([p for p in st['positions'].values() if abs(p)>1e-6])} positions")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["backtest", "regime", "paper", "status"])
    a = ap.parse_args()
    {"backtest": backtest, "regime": regime, "paper": paper, "status": status}[a.mode](CFG)