#!/usr/bin/env python3
"""Small/mid-cap US equity cross-sectional factor gate (DBEQ daily, realistic small-cap costs).

The less-efficient corner: exclude mega-caps (too efficient) and illiquid micro-caps
(untradeable); keep the small/mid liquid band where your small capital is an advantage.
Test XS momentum (12-1 style), short-term reversal, low-vol, residual momentum — GROSS and
NET of ILLIQUIDITY-SCALED cost (small-cap spreads ~30-150bp, the honest killer). Weekly
rebalance + smoothing. Point-in-time universe (incl delisted) -> survivorship-aware.
"""
import math
import os
import sys

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 torch  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"
DAY_NS = 86_400 * 10**9
OUT = "data/surfer/dbeq_ohlcv1d.dbn"
EXCLUDE_TOP = 50      # drop mega-caps (efficient)
DV_FLOOR = 2e6        # $2M/day min (tradeable)
TOPK = 600            # small/mid liquid band size


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, skip=0):
    out = np.full_like(lc, np.nan)
    if skip:
        out[L + skip:] = lc[L:-skip] - lc[:-(L + skip)]
    else:
        out[L:] = lc[L:] - lc[:-L]
    return out


def load():
    import databento as db
    a = db.DBNStore.from_file(OUT).to_ndarray()
    iid = a["instrument_id"]; ts = a["ts_event"].astype(np.int64)
    close = a["close"].astype(np.float64) / 1e9
    vol = a["volume"].astype(np.float64)
    day = ts // DAY_NS
    days = np.unique(day); insts = np.unique(iid)
    dix = np.searchsorted(days, day); iix = np.searchsorted(insts, iid)
    T, N = len(days), len(insts)
    C = np.full((T, N), np.nan); DVOL = np.full((T, N), np.nan)
    C[dix, iix] = np.where(close > 0, close, np.nan)
    DVOL[dix, iix] = close * vol
    return insts, days, C, DVOL


def main():
    insts, days, close, dvol = 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)
    dv30 = roll(np.mean, np.nan_to_num(dvol), 30)
    vol63 = roll(np.std, R, 63)
    year = (1970 + days / 365.25).astype(int)
    print(f"loaded {N} instruments, {T} days ({days.min()}..{days.max()})")

    # point-in-time small/mid-cap universe: drop top mega-caps, require liquidity floor, take band
    univ = np.zeros((T, N), bool)
    for t in range(T):
        elig = np.where((dv30[t] > DV_FLOOR) & np.isfinite(close[t]))[0]
        if len(elig) > EXCLUDE_TOP + 20:
            order = elig[np.argsort(-dv30[t, elig])]
            band = order[EXCLUDE_TOP:EXCLUDE_TOP + TOPK]       # skip mega-caps, take next TOPK
            univ[t, band] = True

    # illiquidity-scaled round-trip cost (bp): small-caps 30-150bp, mid 5-30bp
    rt_cost = np.clip(60.0 / np.sqrt(np.maximum(dv30, 1.0) / 1e6), 5.0, 150.0) / 1e4

    def held_weekly(wt, K=5):
        wh = wt.copy()
        last = 0
        for t in range(T):
            if t % K == 0:
                last = t
            wh[t] = wt[last]
        return wh

    def pnl(sig, net=True):
        s = sig.copy(); s[~univ] = np.nan
        w = xs_weights(s)
        a = 2.0 / (5 + 1)                                       # smooth span 5
        for t in range(1, T):
            w[t] = a * w[t] + (1 - a) * w[t - 1]
        w = held_weekly(w)
        gross = np.sum(w[:-1] * R[1:], axis=1)
        if not net:
            return gross
        turn = np.abs(w[1:] - w[:-1])
        cost = np.sum(turn * rt_cost[1:], axis=1)
        return gross - cost

    sigs = {
        "mom_63_skip5": trailing(lc, 63, skip=5),
        "mom_126_skip5": trailing(lc, 126, skip=5),
        "reversal_5": -trailing(lc, 5),
        "lowvol_63": -vol63,
    }
    print(f"\n===== SMALL/MID-CAP EQUITY FACTOR GATE (band {EXCLUDE_TOP}-{EXCLUDE_TOP+TOPK}, ${DV_FLOOR/1e6:.0f}M floor) =====")
    print(f"universe/day ~{int(univ.sum(1).mean())}, weekly rebal+smooth, deflate N=20")
    print(f"{'factor':>16} {'gross':>6} {'NET':>6} {'IS':>6} {'OOS':>6} {'CPCVmed':>8} {'DSR':>5} | per-year")
    T_ = lambda x: torch.tensor(x[np.isfinite(x)], device=DEV, dtype=torch.float64)
    for nm, sg in sigs.items():
        g = pnl(sg, net=False); p = pnl(sg, net=True)
        gsr = sharpe_t(T_(g)); v = validate(p, days, 20)
        py = " ".join(f"{y}:{sharpe_t(T_(p[year[1:]==y])):+.1f}" for y in range(2023, 2027) if (year[1:] == y).sum() > 40)
        print(f"{nm:>16} {gsr:>+6.2f} {v['full']:>+6.2f} {v['is_']:>+6.2f} {v['oos']:>+6.2f} {v['med']:>+8.2f} {v['dsr']:>5.2f} | {py}")
    print("\nVERDICT: a factor with NET full+OOS+CPCVmed>0 & DSR>0.5 survives small-cap costs = real.")
    print("gross>>NET means the edge is eaten by illiquidity cost (the usual small-cap fate).")


if __name__ == "__main__":
    main()