momentum-system/judge/engine.py

"""Judgement-model engine: signals, features and data windows. This is the single source shared by the
research simulator (data/ticksim) and the live paper trader (judge/live.py).

It is a verbatim port of data/ticksim/sim.py (signals, execution checks) and data/ticksim/ml_dataset.py
(features) as validated by PREREG_v2 on 2026-09-28. Changing anything here changes the model's inputs, so
retrain and re-validate after any edit.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

TICK = 0.01
# the execution config the v2 model was trained and tested under (ml_dataset.EXEC)
EXEC = dict(setup="burst", lat_ms=300, max_spread=0.03, max_chase=0.01, lookback_s=120, vol_mult=4.0,
            burst_pct=0.02, stop_lookback_s=30, min_stop_pct=0.01, max_stop_pct=0.08, run_pct=0.05, run_s=60,
            max_pb=3, lvl_stop_pct=0.015, target_r=0.0, trail_after_r=1.0, trail_pct=0.06, hold_s=1200,
            cooldown_s=30, max_trades=6, price_min=1.0, price_max=20.0)
SETUPS = ("burst", "micropb", "dollar")
FEATS = ["ret_10s", "ret_30s", "ret_60s", "ret_300s", "vol_surge_10s", "vol_surge_60s", "trades_10s", "trades_60s",
         "rv_60s", "off_hod", "vs_vwap", "gain_prev", "up_from_5m_low", "spread", "imb", "ask_size_usd", "min_since_4am",
         "rth", "price", "cumvol_vs_med", "gap_known", "setup_i"]
SETUP_I = {"burst": 0, "micropb": 1, "dollar": 2}


# ----------------------------------------------------------------------------- data windows
def windows(session_date: str, first20_ts: pd.Timestamp | None) -> list[list[pd.Timestamp]]:
    """The tick windows the model was trained on: 09:30-10:30, plus [first +20% minute - 2 min, +60 min]."""
    w = [(pd.Timestamp(f"{session_date} 09:30"), pd.Timestamp(f"{session_date} 10:30"))]
    if first20_ts is not None:
        w.append((max(first20_ts - pd.Timedelta(minutes=2), pd.Timestamp(f"{session_date} 04:00")),
                  first20_ts + pd.Timedelta(minutes=60)))
    w.sort()
    merged = [list(w[0])]
    for s, e in w[1:]:
        if s <= merged[-1][1]:
            merged[-1][1] = max(merged[-1][1], e)
        else:
            merged.append([s, e])
    return merged


def in_windows(ts, ws) -> bool:
    return any(s <= ts < e for s, e in ws)


# ----------------------------------------------------------------------------- bars + signals
def bars5(T: pd.DataFrame) -> pd.DataFrame:
    b = T.set_index("t").resample("5s").agg(o=("p", "first"), h=("p", "max"), l=("p", "min"), c=("p", "last"),
                                            v=("s", "sum"))
    b = b.dropna(subset=["c"])
    b["tc"] = b.index + pd.Timedelta(seconds=5)
    return b.reset_index()


def signals(b: pd.DataFrame, P: dict) -> list[tuple]:
    """[(t, signal_price, stop[, 'cross'])]. burst/dollar: t = the bar close. micropb: t = the breakout bar START
    (resolve the exact crossing trade with cross_time)."""
    out = []
    n = len(b)
    if n < 30:
        return out
    c = b.c.values; h = b.h.values; l = b.l.values; o = b.o.values; v = b.v.values.astype(float)
    t = b.t.values; tc = b.tc.values
    med = pd.Series(v).rolling(24, min_periods=6).median().shift(1).values
    lb = max(1, P["lookback_s"] // 5)
    if P["setup"] == "burst":
        prior_hi = pd.Series(h).rolling(lb, min_periods=1).max().shift(1).values
        c30 = pd.Series(c).shift(6).values
        for i in range(24, n):
            if (h[i] > prior_hi[i] and med[i] > 0 and v[i] >= P["vol_mult"] * med[i] and c30[i] > 0
                    and c[i] / c30[i] - 1 >= P["burst_pct"]):
                k = max(0, i - P["stop_lookback_s"] // 5)
                out.append((tc[i], c[i], float(l[k:i + 1].min())))
    elif P["setup"] == "micropb":
        rs = max(1, P["run_s"] // 5)
        i = rs
        while i < n - 2:
            lo = l[i - rs:i + 1].min(); hi = h[i - rs:i + 1].max()
            if lo > 0 and hi / lo - 1 >= P["run_pct"] and h[i] == hi:
                j = i + 1; pb_low = None
                while j < n and j - i <= P["max_pb"]:
                    if h[j] > hi:
                        break
                    if c[j] < o[j] or c[j] < c[j - 1]:
                        pb_low = l[j] if pb_low is None else min(pb_low, l[j])
                        if hi - pb_low > 0.5 * (hi - lo):
                            pb_low = None; break
                        if j + 1 < n and h[j + 1] > h[j]:
                            out.append((t[j + 1], float(h[j] + TICK), float(pb_low - TICK), "cross"))
                            i = j + 1
                            break
                    j += 1
            i += 1
    elif P["setup"] == "dollar":
        for i in range(24, n):
            for step in (1.0, 0.5):
                lvl = np.floor(c[i] / step) * step
                if lvl > 0 and c[i - 1] < lvl <= h[i] and med[i] > 0 and v[i] >= P["vol_mult"] * med[i]:
                    out.append((tc[i], c[i], lvl * (1 - P["lvl_stop_pct"])))
                    break
    return out


def cross_time(tt: np.ndarray, tp: np.ndarray, bar_start, level: float):
    """The first trade at/above `level` inside the 5-s bar starting at bar_start (no look-ahead)."""
    k0 = np.searchsorted(tt, bar_start, side="left")
    k1 = np.searchsorted(tt, bar_start + np.timedelta64(5, "s"), side="left")
    hit = np.nonzero(tp[k0:k1] >= level)[0]
    return tt[k0 + hit[0]] if len(hit) else None


def entry_check(bid: float, ask: float, sp: float, stop: float, P: dict = EXEC):
    """Execution gates from the simulator. Returns (entry, stop) or None."""
    if bid <= 0 or ask <= 0 or ask < bid:
        return None
    mid = (bid + ask) / 2
    if (ask - bid) / mid > P["max_spread"] or ask > sp * (1 + P["max_chase"]):
        return None
    if not (P["price_min"] <= sp <= P["price_max"]):
        return None
    entry = ask
    stop = min(stop, entry * (1 - P["min_stop_pct"]))
    if (entry - stop) / entry > P["max_stop_pct"] or stop >= entry:
        return None
    return entry, stop


# ----------------------------------------------------------------------------- features
def feats_at(T: pd.DataFrame, Q: pd.DataFrame, ts, prev_close: float) -> dict | None:
    """Features from trades/quotes strictly before ts (verbatim from ml_dataset.feats_at)."""
    tt = T.t.values; k = np.searchsorted(tt, ts, side="left")
    if k < 20:
        return None
    p = T.p.values[:k].astype(float); s = T.s.values[:k].astype(float); t = tt[:k]
    last = p[-1]

    def back(sec):
        return np.searchsorted(t, ts - np.timedelta64(sec, "s"), side="left")

    f = {}
    for sec in (10, 30, 60, 300):
        j = back(sec)
        f[f"ret_{sec}s"] = last / p[j] - 1 if j < k and p[j] > 0 else 0.0
        f[f"vol_{sec}s"] = s[j:].sum()
        f[f"trades_{sec}s"] = k - j
    j120 = back(120); j300 = back(300)
    base = s[j300:j120].sum() / 180.0 if j120 > j300 else np.nan
    f["vol_surge_10s"] = (f["vol_10s"] / 10.0) / base if base and base > 0 else np.nan
    f["vol_surge_60s"] = (f["vol_60s"] / 60.0) / base if base and base > 0 else np.nan
    j60 = back(60)
    lr = np.diff(np.log(p[j60:])) if k - j60 > 3 else np.array([0.0])
    f["rv_60s"] = float(np.std(lr)) * np.sqrt(max(1, len(lr)))
    hod = p.max(); f["off_hod"] = last / hod - 1
    vw = (p * s).sum() / s.sum() if s.sum() else last; f["vs_vwap"] = last / vw - 1
    f["gain_prev"] = last / prev_close - 1
    lo300 = p[j300:].min() if k > j300 else last; f["up_from_5m_low"] = last / lo300 - 1
    qt = Q.t.values; qi = np.searchsorted(qt, ts, side="right") - 1
    if qi >= 0:
        bp, ap = float(Q.bp.iat[qi]), float(Q.ap.iat[qi]); bs, as_ = float(Q.bs.iat[qi]), float(Q["as"].iat[qi])
        mid = (bp + ap) / 2
        f["spread"] = (ap - bp) / mid if mid > 0 else np.nan
        f["imb"] = bs / (bs + as_) if bs + as_ > 0 else np.nan
        f["ask_size_usd"] = as_ * ap
    tsn = pd.Timestamp(ts)
    f["min_since_4am"] = (tsn - tsn.normalize() - pd.Timedelta(hours=4)).total_seconds() / 60
    f["rth"] = float(tsn.time() >= pd.Timestamp("09:30").time())
    f["price"] = last
    f["cum_vol"] = s.sum()
    return f


def finish_features(f: dict, setup: str, gap: float | None, med20: float | None) -> dict:
    """The derived features (ml_fresh.prep): setup code, gap only after 09:30, cumulative volume vs the 20-day median."""
    f = dict(f)
    f["setup_i"] = SETUP_I[setup]
    f["gap_known"] = gap if (f.get("rth") == 1.0 and gap is not None) else np.nan
    f["cumvol_vs_med"] = f["cum_vol"] / med20 if med20 else np.nan
    return f