#!/usr/bin/env python3
"""Measure best-level Order-Flow-Imbalance (OFI) predictive edge on ES MBP-10.

Signal-validation harness (NO training, NO RL). Answers: does OFI predict
forward mid-returns out-of-sample, and is the predicted move ever larger than
the spread we pay to CROSS? Per the 2026-06-05 omnisearch synthesis
(pearl_omnisearch_synthesis_signal_first_and_mtm_reward): prove edge exists
before any more RL/reward work.

Method (Cont-Kukanov-Stoikov 2014 OFI + standard IC validation):
  * Decode clean MBP-10 via databento_dbn (reads the authoritative
    mbp10.levels[0] directly — the L0 decode bug was in the Rust parser, not
    the data, so this measures the TRUE achievable edge).
  * Front-month filter (dominant instrument_id).
  * Event-level OFI; mid = (bid+ask)/2.
  * Bucket by wall-clock time; per-bucket sum(OFI), bucket-close mid.
  * Contemporaneous regression Δmid ~ OFI  → R² (sanity: clean book ⇒ high R²;
    a broken book ⇒ low R², which would have caught the L0 bug).
  * Predictive: corr(OFI_t, forward_return_{t→t+h}) across horizons h.
    Reports Pearson IC, Spearman IC, sign-accuracy, Newey-West t-stat.
  * OOS: 70/30 time split (fit/test) + run on two quarters to compare stability.
  * CROSSING VERDICT (headline): predicted move in TICKS for large-|OFI|
    signals vs the 1-tick ($12.50) round-trip cost of crossing the spread.

Usage:
  python3 scripts/measure_ofi_ic.py <file.dbn.zst> [--n-cap N] [--bucket-sec S]
  python3 scripts/measure_ofi_ic.py --self-test
"""
import argparse
import sys
import time

import numpy as np

TICK = 0.25          # ES tick size in index points
TICK_USD = 12.50     # $ per tick per contract
PX_SCALE = 1e9       # dbn fixed-point price scale (1 unit = 1e-9)


def compute_ofi(bid_px, bid_sz, ask_px, ask_sz):
    """Best-level OFI per consecutive event (Cont-Kukanov-Stoikov 2014).

    OFI_n = [q^b_n·1{P^b_n>=P^b_{n-1}} - q^b_{n-1}·1{P^b_n<=P^b_{n-1}}]
          + [q^a_{n-1}·1{P^a_n>=P^a_{n-1}} - q^a_n·1{P^a_n<=P^a_{n-1}}]
    Bid liquidity growth ⇒ OFI up (buy pressure); ask growth ⇒ OFI down.
    Returns array length len-1 (aligned to event n=1..len-1).
    """
    pb_n, pb_p = bid_px[1:], bid_px[:-1]
    qb_n, qb_p = bid_sz[1:].astype(np.float64), bid_sz[:-1].astype(np.float64)
    pa_n, pa_p = ask_px[1:], ask_px[:-1]
    qa_n, qa_p = ask_sz[1:].astype(np.float64), ask_sz[:-1].astype(np.float64)
    e_bid = qb_n * (pb_n >= pb_p) - qb_p * (pb_n <= pb_p)
    e_ask = qa_p * (pa_n >= pa_p) - qa_n * (pa_n <= pa_p)
    return e_bid + e_ask


def self_test():
    # bid_px, bid_sz, ask_px, ask_sz
    bid_px = np.array([100, 100, 101], dtype=np.int64)
    bid_sz = np.array([10, 15, 8], dtype=np.int64)
    ask_px = np.array([101, 101, 102], dtype=np.int64)
    ask_sz = np.array([10, 10, 12], dtype=np.int64)
    ofi = compute_ofi(bid_px, bid_sz, ask_px, ask_sz)
    # e1: bid same +5, ask same 0 -> 5 ; e2: bid up +8, ask up +10 -> 18
    expected = np.array([5.0, 18.0])
    assert np.allclose(ofi, expected), f"OFI self-test FAILED: {ofi} != {expected}"
    print("OFI self-test PASS:", ofi.tolist())
    return 0


def newey_west_tstat(x, y, lags):
    """t-stat of slope b in y = a + b·x with Newey-West HAC standard error."""
    n = len(x)
    if n < lags + 3:
        return float("nan")
    X = np.column_stack([np.ones(n), x])
    xtx_inv = np.linalg.inv(X.T @ X)
    beta = xtx_inv @ (X.T @ y)
    resid = y - X @ beta
    # HAC meat matrix S = Σ_l w_l Σ_t u_t u_{t-l} x_t x_{t-l}'
    u = resid[:, None] * X  # [n,2]
    S = u.T @ u
    for l in range(1, lags + 1):
        w = 1.0 - l / (lags + 1.0)
        G = u[l:].T @ u[:-l]
        S += w * (G + G.T)
    cov = xtx_inv @ S @ xtx_inv
    se_b = np.sqrt(cov[1, 1])
    return float(beta[1] / se_b) if se_b > 0 else float("nan")


def decode(path, n_cap):
    import zstandard
    import databento_dbn as d
    ts = np.zeros(n_cap, np.int64)
    bpx = np.zeros(n_cap, np.int64); apx = np.zeros(n_cap, np.int64)
    bsz = np.zeros(n_cap, np.int64); asz = np.zeros(n_cap, np.int64)
    inst = np.zeros(n_cap, np.int32)
    dec = d.DBNDecoder(); dctx = zstandard.ZstdDecompressor()
    i = 0; t0 = time.time(); done = False
    with open(path, "rb") as fh:
        reader = dctx.stream_reader(fh)
        while not done:
            chunk = reader.read(1 << 22)
            if not chunk:
                break
            dec.write(chunk)
            for rec in dec.decode():
                if "MBP10" not in type(rec).__name__:
                    continue
                l0 = rec.levels[0]
                ts[i] = rec.ts_event
                bpx[i] = l0.bid_px; apx[i] = l0.ask_px
                bsz[i] = l0.bid_sz; asz[i] = l0.ask_sz
                inst[i] = rec.instrument_id
                i += 1
                if i >= n_cap:
                    done = True
                    break
    dt = time.time() - t0
    return (ts[:i], bpx[:i], apx[:i], bsz[:i], asz[:i], inst[:i], dt)


def pct(x):
    return "n/a" if x is None or (isinstance(x, float) and np.isnan(x)) else f"{x:.4f}"


def analyze(path, n_cap, bucket_sec):
    ts, bpx, apx, bsz, asz, inst, dt = decode(path, n_cap)
    n = len(ts)
    ids, cnts = np.unique(inst, return_counts=True)
    dom = ids[cnts.argmax()]
    m = inst == dom
    ts, bpx, apx, bsz, asz = ts[m], bpx[m], apx[m], bsz[m], asz[m]
    print(f"\n=== {path.split('/')[-1]} ===")
    print(f"decoded {n} MBP10 events in {dt:.1f}s | instruments={len(ids)} "
          f"dominant={dom} ({100*cnts.max()/n:.1f}%) → {m.sum()} front-month events")

    # valid uncrossed L0 only
    valid = (bpx > 0) & (apx > 0) & (bpx < apx)
    crossed = (~valid).mean()
    ts, bpx, apx, bsz, asz = ts[valid], bpx[valid], apx[valid], bsz[valid], asz[valid]
    mid = (bpx + apx) / 2.0 / PX_SCALE
    spread_ticks = ((apx - bpx) / PX_SCALE / TICK)
    print(f"crossed/invalid L0 dropped: {100*crossed:.3f}% | "
          f"spread ticks p50={np.median(spread_ticks):.2f} p90={np.percentile(spread_ticks,90):.2f}")

    ofi = compute_ofi(bpx, bsz, apx, asz)        # len-1, aligned to events 1..
    ts_e, mid_e = ts[1:], mid[1:]                # event mids aligned to ofi

    # bucket by wall-clock
    bsz_ns = int(bucket_sec * 1e9)
    bk = (ts_e - ts_e[0]) // bsz_ns
    # np.unique returns (unique, index, inverse) in that fixed order.
    ub, first, inv = np.unique(bk, return_index=True, return_inverse=True)
    inv = np.asarray(inv).ravel()                              # numpy 2.x shape guard
    ofi_b = np.bincount(inv, weights=ofi)                       # sum OFI per bucket
    last = np.append(first[1:] - 1, len(bk) - 1)
    close_mid = mid_e[last]                                     # bucket-close mid
    bk_id = ub                                                  # integer bucket index (gaps possible)
    nb = len(ub)
    print(f"buckets({bucket_sec}s)={nb} | front-month span "
          f"≈{(ts_e[-1]-ts_e[0])/1e9/3600:.1f}h")

    # contemporaneous: Δmid over bucket ~ OFI_bucket  (R² sanity)
    dmid = np.empty(nb); dmid[0] = 0.0
    dmid[1:] = (close_mid[1:] - close_mid[:-1])
    contig1 = np.empty(nb, bool); contig1[0] = False
    contig1[1:] = (bk_id[1:] - bk_id[:-1]) == 1
    cx, cy = ofi_b[contig1], dmid[contig1]
    if len(cx) > 10 and cx.std() > 0:
        from scipy.stats import pearsonr
        r = pearsonr(cx, cy)[0]
        b = np.cov(cx, cy)[0, 1] / np.var(cx)
        print(f"\nCONTEMPORANEOUS Δmid~OFI : R²={r**2:.4f}  "
              f"beta={b:.3e} pts/OFI  (sanity: clean book ⇒ R² high)")

    # predictive IC across horizons (contiguous forward windows only)
    from scipy.stats import pearsonr, spearmanr
    horizons = [1, 2, 5, 10, 30, 60]
    print(f"\nPREDICTIVE  corr(OFI_t, fwd_return_{{t→t+h}})   bucket={bucket_sec}s")
    print(f"{'h(s)':>5} {'n':>8} {'IC_pear':>9} {'IC_spear':>9} {'sign%':>7} "
          f"{'NW_t':>7} {'beta(ticks/OFI)':>16} {'OOS_IC':>8}")
    results = {}
    for h in horizons:
        # forward log-return over exactly-contiguous h buckets
        ok = (bk_id[h:] - bk_id[:-h]) == h
        x = ofi_b[:-h][ok]
        fwd = np.log(close_mid[h:][ok] / close_mid[:-h][ok])
        nz = x != 0
        x, fwd = x[nz], fwd[nz]
        if len(x) < 100 or x.std() == 0 or fwd.std() == 0:
            continue
        ic = pearsonr(x, fwd)[0]
        ics = spearmanr(x, fwd)[0]
        sign = np.mean(np.sign(x) == np.sign(fwd))
        nwt = newey_west_tstat(x, fwd, lags=h)
        beta = np.cov(x, fwd)[0, 1] / np.var(x)        # log-return per unit OFI
        beta_ticks = beta * np.median(close_mid) / TICK  # ticks per unit OFI
        # OOS: fit not needed for IC; measure IC on last 30%
        k = int(0.7 * len(x))
        oos_ic = pearsonr(x[k:], fwd[k:])[0] if len(x) - k > 50 else float("nan")
        results[h] = dict(n=len(x), ic=ic, ics=ics, sign=sign, nwt=nwt,
                          beta_ticks=beta_ticks, oos=oos_ic)
        print(f"{h:>5} {len(x):>8} {ic:>9.4f} {ics:>9.4f} {100*sign:>6.2f}% "
              f"{nwt:>7.2f} {beta_ticks:>16.3e} {pct(oos_ic):>8}")

    # CROSSING economics: predicted move in ticks at large |OFI| vs 1-tick cost
    print(f"\nCROSSING VERDICT (cost = 1 tick round-trip = ${TICK_USD*2:.2f}):")
    for h in [1, 5, 10]:
        if h not in results:
            continue
        ok = (bk_id[h:] - bk_id[:-h]) == h
        x = ofi_b[:-h][ok]
        fwd = np.log(close_mid[h:][ok] / close_mid[:-h][ok])
        nz = x != 0; x, fwd = x[nz], fwd[nz]
        if len(x) < 100:
            continue
        beta = np.cov(x, fwd)[0, 1] / np.var(x)
        thr = np.percentile(np.abs(x), 90)            # trade only on top-decile |OFI|
        big = np.abs(x) >= thr
        pred_move_ticks = np.abs(beta * x[big]) * np.median(close_mid) / TICK
        print(f"  h={h}s: top-decile |OFI| trades → predicted |move| ticks "
              f"p50={np.median(pred_move_ticks):.3f} p90={np.percentile(pred_move_ticks,90):.3f} "
              f"max={pred_move_ticks.max():.3f} | frac>1tick={100*np.mean(pred_move_ticks>1.0):.2f}%")
    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("file", nargs="?", help="path to .dbn.zst MBP-10 file")
    ap.add_argument("--n-cap", type=int, default=8_000_000)
    ap.add_argument("--bucket-sec", type=float, default=1.0)
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args()
    if args.self_test:
        return self_test()
    if not args.file:
        ap.error("file required (or --self-test)")
    analyze(args.file, args.n_cap, args.bucket_sec)
    return 0


if __name__ == "__main__":
    sys.exit(main())