#!/usr/bin/env python3
"""Surfer MFT quality test — regime-adaptive, selective, intraday (GPU).

Tests the ACTUAL surfer thesis (not daily CTA): ~5-min decisions, ~15-min holds,
flat-by-close, regime-ADAPTIVE (ride trend / fade range / stand aside in chop),
and QUALITY over quantity (only the top-conviction waves). Intraday + flat-by-close
⇒ no overnight, no roll-gap problem.

Data: clean ES OHLCV-1m (test_data/futures-baseline/ES.FUT, front-month per quarter),
resampled to 5-min. Signal/regime/backtest/stats on the GPU (torch.cuda); load=CPU staging.

Pre-registered (NOT tuned), to bound overfit on a selective/small-sample test:
  vol      = trailing 12-bar (1h) std of 5-min returns
  trend t  = ret(12 bars) / (vol·√12)              # 1h trend t-stat
  TREND  if |t| > 1.0  → trade WITH ret(12)        # ride the wave
  RANGE  if |t| < 0.5  → trade AGAINST ret(3)       # fade small chop toward mean
  else (0.5..1.0)      → STAND ASIDE (no trade)
  conviction = |t| (trend) or |ret3|/(vol·√3) (range)
  entry=close[t], exit=close[t+3] (15 min), cost = 1.0 tick round-trip (crossing).
QUALITY curve: rank candidate trades by conviction, take the top q%, report per-trade
net edge (ticks), hit-rate, count — IS(first 70%) and OOS(last 30%).
"""
import glob
import math

import numpy as np
import torch

DEV = "cuda" if torch.cuda.is_available() else None
TICK = 0.25
COST_TICKS = 1.0
BAR_NS = 5 * 60 * 10**9
DAY_NS = 86_400 * 10**9


def load_es_5min():
    import databento as db
    # Prefer the full-history continuous front-month pull (data/surfer/es1m/), else fall back
    # to the local ~2y parent OHLCV-1m (futures-baseline/ES, per-quarter front-month pick).
    es1m = sorted(glob.glob("data/surfer/es1m/*.dbn"))
    ts_all, c_all = [], []
    if es1m:
        for p in es1m:                                            # to_ndarray = light (to_df OOMs at 20GB)
            try:
                arr = db.DBNStore.from_file(p).to_ndarray()
            except Exception:
                continue
            if len(arr) == 0 or "close" not in arr.dtype.names:
                continue
            ts = arr["ts_event"].astype(np.int64)
            c = arr["close"].astype(np.float64) / 1e9            # raw 1e9 fixed-point → price
            m = c > 0
            ts_all.append(ts[m]); c_all.append(c[m])
    else:
        for p in sorted(glob.glob("test_data/futures-baseline/ES.FUT/*.dbn.zst")):
            try:
                df = db.DBNStore.from_file(p).to_df().reset_index()
            except Exception:
                continue
            if df.empty or "close" not in df.columns:
                continue
            df = df[df["close"] > 0]
            if df.empty:
                continue
            dom = df["instrument_id"].value_counts().idxmax()    # front month for the quarter
            df = df[df["instrument_id"] == dom].sort_values("ts_event")
            ts_all.append(df["ts_event"].astype("int64").to_numpy())
            c_all.append(df["close"].to_numpy(np.float64))
    ts = np.concatenate(ts_all); c = np.concatenate(c_all)
    _u, _ui = np.unique(ts, return_index=True)                   # dedup any overlapping ts at chunk seams
    ts, c = ts[_ui], c[_ui]
    o = np.argsort(ts, kind="stable"); ts, c = ts[o], c[o]
    b5 = ts // BAR_NS                                            # 5-min bin id
    _, first = np.unique(b5, return_index=True)
    last = np.append(first[1:] - 1, len(b5) - 1)                 # last 1-min close in each 5-min bin
    return ts[last], c[last]


def rolling_std(x, w):  # trailing-w std at each t (NaN first w-1); O(T) memory (cumsum, no unfold)
    z = torch.zeros(1, dtype=x.dtype, device=x.device)
    cs = torch.cat([z, torch.cumsum(x, 0)])
    cs2 = torch.cat([z, torch.cumsum(x * x, 0)])
    s = cs[w:] - cs[:-w]
    s2 = cs2[w:] - cs2[:-w]
    mean = s / w
    var = (s2 / w - mean * mean).clamp_min(0.0)
    out = torch.full_like(x, float("nan"))
    out[w - 1:] = var.sqrt()
    return out


def run():
    ts5, close5 = load_es_5min()
    day = ts5 // DAY_NS
    lc = torch.tensor(np.log(close5), device=DEV, dtype=torch.float64)
    px = torch.tensor(close5, device=DEV, dtype=torch.float64)
    dy = torch.tensor(day, device=DEV)
    T = len(lc)
    t = torch.arange(T, device=DEV)

    def sameday(lag):  # day[i]==day[i-lag]
        m = torch.zeros(T, dtype=torch.bool, device=DEV)
        m[lag:] = dy[lag:] == dy[:-lag]
        return m

    r1 = torch.zeros_like(lc); r1[1:] = lc[1:] - lc[:-1]
    r1 = torch.where(sameday(1), r1, torch.zeros_like(r1))       # zero overnight returns
    vol = rolling_std(r1, 12) * math.sqrt(1.0)                   # per-bar vol (5-min)
    ret12 = torch.full_like(lc, float("nan")); ret12[12:] = lc[12:] - lc[:-12]
    ret3 = torch.full_like(lc, float("nan")); ret3[3:] = lc[3:] - lc[:-3]
    tstat = ret12 / (vol * math.sqrt(12)).clamp_min(1e-9)

    # forward 15-min (3 bars) PRICE move in ticks, intraday only
    fwd_ticks = torch.full_like(px, float("nan"))
    fwd_ticks[:-3] = (px[3:] - px[:-3]) / TICK
    fwd_valid = torch.zeros(T, dtype=torch.bool, device=DEV)
    fwd_valid[:-3] = dy[3:] == dy[:-3]

    valid = fwd_valid & sameday(12) & torch.isfinite(tstat) & torch.isfinite(vol) & (vol > 0)

    def evaluate(name, direction, conviction, candidate):
        cand = candidate & valid & torch.isfinite(direction) & (direction != 0)
        # net per-trade edge in ticks (enter close[t], exit close[t+3], 1-tick round-trip cost)
        net = direction * fwd_ticks - COST_TICKS
        idx = torch.nonzero(cand, as_tuple=True)[0]
        if len(idx) < 50:
            print(f"  [{name}] too few candidates ({len(idx)})"); return
        conv = conviction[idx]; netc = net[idx]; tt = t[idx]
        split = int(0.7 * T)
        is_m = tt < split; oos_m = ~is_m
        print(f"  [{name}] candidates={len(idx)}  (IS {int(is_m.sum())} / OOS {int(oos_m.sum())})")
        print(f"      {'top-q':>6} {'n_OOS':>6} {'net_t/trade_OOS':>15} {'t-stat':>7} {'hit%_OOS':>9} {'IS':>8}")
        order = torch.argsort(conv, descending=True)
        for q in [0.05, 0.10, 0.25, 0.50, 1.00]:
            k = max(int(q * len(idx)), 1)
            sel = order[:k]
            sm = oos_m[sel]; im = is_m[sel]
            no = int(sm.sum())
            if no < 10:
                print(f"      {q:>6.2f}  (OOS n<10)"); continue
            vals = netc[sel][sm]
            net_oos = float(vals.mean())
            tstat = net_oos / (float(vals.std()) / math.sqrt(no) + 1e-12)
            hit_oos = float((vals > 0).float().mean())
            net_is = float(netc[sel][im].mean()) if int(im.sum()) > 0 else float("nan")
            print(f"      {q:>6.2f} {no:>6} {net_oos:>+15.3f} {tstat:>+7.2f} {100*hit_oos:>8.1f}% {net_is:>+8.3f}")

    print("\n================ SURFER MFT QUALITY TEST (GPU, intraday flat-by-close) ================")
    print(f"device={DEV}  5-min bars={T}  span days≈{int((day.max()-day.min()))}  cost={COST_TICKS} tick round-trip")
    print("net_t/trade = mean per-trade edge in TICKS after cost; quality bar: OOS top-decile > 0 & hit>50%\n")

    # regime-adaptive direction + conviction
    trend = tstat.abs() > 1.0
    rang = tstat.abs() < 0.5
    dir_adapt = torch.zeros_like(tstat)
    dir_adapt = torch.where(trend, torch.sign(ret12), dir_adapt)        # ride
    dir_adapt = torch.where(rang, -torch.sign(ret3), dir_adapt)         # fade
    conv_adapt = torch.where(trend, tstat.abs(),
                  torch.where(rang, (ret3.abs() / (vol * math.sqrt(3)).clamp_min(1e-9)),
                              torch.zeros_like(tstat)))
    evaluate("REGIME-ADAPTIVE (ride trend / fade range / aside)", dir_adapt, conv_adapt, (trend | rang))

    # baselines for contrast
    evaluate("naive MOMENTUM always", torch.sign(ret12), tstat.abs(), torch.ones(T, dtype=torch.bool, device=DEV))
    evaluate("naive MEAN-REVERSION always", -torch.sign(ret3), (ret3.abs()/(vol*math.sqrt(3)).clamp_min(1e-9)),
             torch.ones(T, dtype=torch.bool, device=DEV))


def main():
    if DEV is None:
        raise SystemExit("CUDA not available")
    run()


if __name__ == "__main__":
    main()