momentum-system/research/regime.py

"""research/regime.py — Plan 011: a FREE-data daily signal that decides
WHETHER the ORB trades at all today (a session-level metalabel), orthogonal to
WHICH NAME to trade (strategy/population.py + rank_one_bullet). Predicts
``P(our ORB makes money THIS session)`` from pre-open features; the label is
our OWN per-session realized expR (Chan metalabeling) — the mean of
``research.experiments.replay._sim_baseline_r`` (the LIVE exit config
simulated on the bar path) over that session's setup fires.

  python -m research.regime probe                        # Step 0 STOP gate
  python -m research.regime build-features                # Step 1: regime_features
  python -m research.regime score --model threshold_v1     # Step 1: walk-forward rule
  python -m research.regime score --model meta_v1          # Step 2: walk-forward GBM

DISCIPLINE (see docs/research/regime_notes.md / plan 011):
  * Every score is WALK-FORWARD: favorable(D) uses ONLY rows with
    session_date < D (1-session embargo). Nothing here may read a full-sample
    statistic — see ``score_threshold``/``score_meta`` and
    ``test_regime_score_is_walk_forward``, the leakage guard.
  * ``label_expr``/``label_pos`` are OUTCOME columns; FEATURE_COLS never
    includes them (``test_regime_features_no_label_leak``).
  * ``full_sample_fit_threshold`` (Step 4) is DELIBERATELY LEAKY — a sanity
    check only, never the shipped flag.

Imports nothing from the live loop (trading_execution.py / run_live_paper.py /
strategy/exits.py's live call sites / strategy/population.py are never
imported here). Reuses research.experiments.replay's fire reconstruction +
_sim_baseline_r as BOTH the fire source and the label source.
"""

from __future__ import annotations

import argparse
import os
from bisect import bisect_left
from collections import defaultdict
from datetime import date, datetime, timezone
from pathlib import Path
from statistics import mean, median, pstdev

_SCHEMA_SQL = Path(__file__).resolve().parent.parent / "storage" / "regime_schema.sql"

# Pre-open / prior-close-lagged features only. See module docstring + plan 011
# Step 1 for the source of each (breadth_* and tape_regime come from the
# session's OWN setup_events rows — the same "as_of=open" contract
# strategy/population.py relies on; spy_* are lagged to the PRIOR close;
# bot_trailing_expr_10d is a rolling mean of PRIOR sessions' labels).
FEATURE_COLS = (
    "breadth_gap", "breadth_rvol", "tape_regime",
    "spy_mom5", "spy_mom20", "spy_vol20", "bot_trailing_expr_10d",
)
LABEL_COLS = ("label_expr", "label_pos")

TAPE_ORDINAL = {"trend_down": -1.0, "chop": 0.0, "trend_up": 1.0}

MIN_LABELABLE_SESSIONS = 150   # Step 0 STOP threshold (thin sample)
SPY_LOOKBACK_DAYS = 20
TRAILING_WINDOW = 10            # bot_trailing_expr_10d
MIN_THRESHOLD_PRIOR = 5         # threshold_v1: below this, can't infer a sign -> fail open
MIN_TRAIN_SESSIONS = 60         # meta_v1: below this, don't gate on an unfit model -> fail open
P_MIN_DEFAULT = 0.5

# Fidelity floors used to simulate the session label (matches the harness's
# default adverse-fill accounting for setup-sourced fires — see replay.py).
_LABEL_FIDELITY = {"atr_frac": 0.25, "spread_mult": 3, "adverse_ticks": 1}


def apply_schema(con) -> None:
    """Create regime_features/regime_scores idempotently. Honors PG_SKIP_SCHEMA
    exactly like research.experiments.registry.apply_schema (owner-provisioned
    DDL escape hatch)."""
    if os.environ.get("PG_SKIP_SCHEMA", "").strip().lower() in {"1", "true", "yes"}:
        return
    con.execute(_SCHEMA_SQL.read_text())


def _as_date(v) -> date:
    return v if isinstance(v, date) else date.fromisoformat(str(v)[:10])


# --------------------------------------------------------------------------- #
# Step 0 — data-availability probe (STOP gate)
# --------------------------------------------------------------------------- #
def _labelable_sessions(con) -> list[date]:
    rows = con.execute(
        "SELECT DISTINCT session_date FROM setup_events "
        "WHERE entry_reference_price IS NOT NULL AND invalidation_price IS NOT NULL "
        "ORDER BY session_date").fetchall()
    return [_as_date(r[0]) for r in rows]


def _spy_daily_dates(con) -> list[date]:
    rows = con.execute(
        "SELECT trade_date FROM daily_bars WHERE symbol='SPY' ORDER BY trade_date").fetchall()
    return [_as_date(r[0]) for r in rows]


def probe(con) -> dict:
    """Step 0 STOP gate. Pure counts + a `stop` verdict — under ~150 labelable
    sessions is a thin sample (STOP). SPY coverage is reported separately
    (never a STOP by itself): when it doesn't cover the setup range, build-
    features falls back to the non-SPY features (spy_* stay NULL) rather than
    fabricating SPY history."""
    sessions = _labelable_sessions(con)
    n_labelable = len(sessions)
    spy_dates = _spy_daily_dates(con)
    n_with_spy20 = sum(1 for d in sessions if bisect_left(spy_dates, d) >= SPY_LOOKBACK_DAYS)
    spy_span = (spy_dates[0], spy_dates[-1]) if spy_dates else (None, None)
    spy_overlaps_range = bool(spy_dates) and bool(sessions) and spy_dates[0] <= sessions[-1]
    return {
        "labelable_sessions": n_labelable,
        "spy_daily_count": len(spy_dates),
        "spy_span": spy_span,
        "sessions_with_spy20": n_with_spy20,
        "spy_coverage_ok": spy_overlaps_range and n_with_spy20 >= 0.5 * max(n_labelable, 1),
        "stop": n_labelable < MIN_LABELABLE_SESSIONS,
    }


# --------------------------------------------------------------------------- #
# Step 1 — build-features
# --------------------------------------------------------------------------- #
class _ConShim:
    """Wraps a raw connection as a store-like object (`.con`) so
    replay.reconstruct_fires_from_setups can be reused verbatim without a full
    EventStore."""

    def __init__(self, con):
        self.con = con


def _session_labels(con, sessions: list[date]) -> dict[date, float]:
    """{session_date: label_expr} = mean(_sim_baseline_r) over that session's
    setup fires — the SAME simulation the exit-experiment harness uses for a
    setup-sourced fire's baseline. Sessions with no computable fire (no bars,
    no risk) are omitted, not zero-filled."""
    from research.experiments.replay import _entry_index, _sim_baseline_r, reconstruct_fires_from_setups
    from research.query import query_minute_bars

    fires = reconstruct_fires_from_setups(_ConShim(con))
    by_session: dict[date, list] = defaultdict(list)
    for f in fires:
        by_session[f.session_date].append(f)

    out: dict[date, float] = {}
    for d in sessions:
        rs = []
        for f in by_session.get(d, []):
            if f.entry_price is None or f.stop_price is None or f.entry_price - f.stop_price <= 0:
                continue
            bars = query_minute_bars(con, f.symbol, f.session_date)
            if bars is None or not len(bars):
                continue
            idx = _entry_index(bars, f.fill_ts)
            bars_after = bars.iloc[idx + 1:].reset_index(drop=True)
            r = _sim_baseline_r(f, bars_after, _LABEL_FIDELITY)
            if r is not None:
                rs.append(r)
        if rs:
            out[d] = mean(rs)
    return out


def _session_context(con, sessions: list[date]) -> dict[date, dict]:
    """{session_date: {breadth_gap, breadth_rvol, tape_regime}} from that
    session's OWN labelable setup_events rows (gap_pct/relative_volume/
    tape_regime_at_trigger — the same open-knowable columns
    strategy/population.py and the labeler already treat as pre-entry
    context). tape_regime is the EARLIEST (by setup_time) setup's reading."""
    rows = con.execute(
        "SELECT session_date, gap_pct, relative_volume, tape_regime_at_trigger, setup_time "
        "FROM setup_events WHERE entry_reference_price IS NOT NULL "
        "AND invalidation_price IS NOT NULL ORDER BY session_date, setup_time").fetchall()
    by_session: dict[date, list] = defaultdict(list)
    for sd, gap, rvol, tape, _st in rows:
        by_session[_as_date(sd)].append((gap, rvol, tape))

    out = {}
    for d in sessions:
        items = by_session.get(d, [])
        if not items:
            out[d] = {"breadth_gap": None, "breadth_rvol": None, "tape_regime": None}
            continue
        gaps = [g for g, _r, _t in items if g is not None]
        rvols = [r for _g, r, _t in items if r is not None]
        out[d] = {
            "breadth_gap": (sum(1 for g in gaps if g > 0) / len(gaps)) if gaps else None,
            "breadth_rvol": median(rvols) if rvols else None,
            "tape_regime": TAPE_ORDINAL.get(items[0][2]),
        }
    return out


def _spy_features(con, sessions: list[date]) -> dict[date, dict]:
    """{session_date: {spy_mom5, spy_mom20, spy_vol20}} strictly lagged to the
    PRIOR close (all SPY daily bars used are trade_date < session_date).
    Missing/insufficient SPY history -> None (never fabricated)."""
    rows = con.execute(
        "SELECT trade_date, close FROM daily_bars WHERE symbol='SPY' ORDER BY trade_date").fetchall()
    dates = [_as_date(r[0]) for r in rows]
    closes = [float(r[1]) for r in rows]

    out = {}
    for d in sessions:
        idx = bisect_left(dates, d)   # count of SPY rows strictly before d
        prior = closes[:idx]
        mom5 = (prior[-1] / prior[-6] - 1) if len(prior) >= 6 else None
        mom20 = (prior[-1] / prior[-21] - 1) if len(prior) >= 21 else None
        vol20 = None
        if len(prior) >= 21:
            rets = [prior[i] / prior[i - 1] - 1 for i in range(-20, 0)]
            vol20 = pstdev(rets)
        out[d] = {"spy_mom5": mom5, "spy_mom20": mom20, "spy_vol20": vol20}
    return out


def _trailing_expr(labels: dict[date, float], sessions: list[date]) -> dict[date, float | None]:
    """mean(label_expr) over the PRIOR (up to TRAILING_WINDOW) labelable
    sessions — strictly backward-looking: `sessions` must be sorted ascending,
    and a session's own label is appended to history AFTER its trailing value
    is read."""
    out: dict[date, float | None] = {}
    hist: list[float] = []
    for d in sessions:
        window = hist[-TRAILING_WINDOW:]
        out[d] = mean(window) if window else None
        if d in labels:
            hist.append(labels[d])
    return out


def build_features(con, since: str | None = None) -> int:
    """Step 1: one regime_features row per labelable session with a computable
    label. Idempotent (PK session_date + ON CONFLICT DO UPDATE — features are
    safe to rebuild as more setup/SPY data lands)."""
    apply_schema(con)
    all_sessions = _labelable_sessions(con)
    labels = _session_labels(con, all_sessions)       # full history: label per session
    trailing = _trailing_expr(labels, all_sessions)   # full history: correct rolling window

    sessions = all_sessions
    if since:
        cutoff = date.fromisoformat(str(since))
        sessions = [d for d in sessions if d >= cutoff]   # only THESE rows get written

    ctx = _session_context(con, sessions)
    spy = _spy_features(con, sessions)

    now = datetime.now(timezone.utc)
    n = 0
    for d in sessions:
        if d not in labels:
            continue   # no computable fire for this session -> unlabelable, skip
        row = {**ctx[d], **spy[d], "bot_trailing_expr_10d": trailing[d],
               "label_expr": labels[d], "label_pos": labels[d] > 0}
        con.execute(
            "INSERT INTO regime_features (session_date, breadth_gap, breadth_rvol, "
            "tape_regime, spy_mom5, spy_mom20, spy_vol20, bot_trailing_expr_10d, "
            "label_expr, label_pos, computed_at) VALUES (?,?,?,?,?,?,?,?,?,?,?) "
            "ON CONFLICT (session_date) DO UPDATE SET "
            "breadth_gap=EXCLUDED.breadth_gap, breadth_rvol=EXCLUDED.breadth_rvol, "
            "tape_regime=EXCLUDED.tape_regime, spy_mom5=EXCLUDED.spy_mom5, "
            "spy_mom20=EXCLUDED.spy_mom20, spy_vol20=EXCLUDED.spy_vol20, "
            "bot_trailing_expr_10d=EXCLUDED.bot_trailing_expr_10d, "
            "label_expr=EXCLUDED.label_expr, label_pos=EXCLUDED.label_pos, "
            "computed_at=EXCLUDED.computed_at",
            [d, row["breadth_gap"], row["breadth_rvol"], row["tape_regime"],
             row["spy_mom5"], row["spy_mom20"], row["spy_vol20"],
             row["bot_trailing_expr_10d"], row["label_expr"], row["label_pos"], now])
        n += 1
    return n


# --------------------------------------------------------------------------- #
# Step 1 — threshold_v1 (no ML, walk-forward)
# --------------------------------------------------------------------------- #
def _tape_favorable_set(prior_rows: list[dict]) -> set[float]:
    """Which tape_regime ordinal values have positive mean label_expr over
    PRIOR rows ONLY — recomputed at every D, never a fixed global set (or it
    would be a full-sample fit)."""
    by_tape: dict[float, list[float]] = defaultdict(list)
    for r in prior_rows:
        if r.get("tape_regime") is not None and r.get("label_expr") is not None:
            by_tape[r["tape_regime"]].append(r["label_expr"])
    return {t for t, rs in by_tape.items() if mean(rs) > 0}


def _vol_side_ok(prior_rows: list[dict], today_vol) -> bool:
    """Whichever SPY-vol side (> median vs <= median, computed over PRIOR rows
    only) has the higher mean label_expr; True if today's spy_vol20 is on that
    side. Fail-open (True) when uncomputable — do NOT hardcode low-vol=good;
    let prior data pick the sign."""
    have = [r for r in prior_rows if r.get("spy_vol20") is not None and r.get("label_expr") is not None]
    if today_vol is None or len(have) < 10:
        return True
    vols = sorted(r["spy_vol20"] for r in have)
    med = vols[len(vols) // 2]
    hi = [r["label_expr"] for r in have if r["spy_vol20"] > med]
    lo = [r["label_expr"] for r in have if r["spy_vol20"] <= med]
    if not hi or not lo:
        return True
    prefer_hi = mean(hi) > mean(lo)
    return (today_vol > med) == prefer_hi


def score_threshold(rows: list[dict]) -> list[dict]:
    """Walk-forward threshold_v1 rule (no fit, cannot leak by construction):
    favorable(D) reads row D's OWN features plus `rows[:i]` (strictly PRIOR
    sessions) only — never anything at index >= i. `rows` must be sorted by
    session_date ascending. Returns [{session_date, p_profit, favorable}, ...].
    """
    out = []
    for i, row in enumerate(rows):
        prior = rows[:i]
        if len(prior) < MIN_THRESHOLD_PRIOR:
            out.append({"session_date": row["session_date"], "p_profit": 1.0, "favorable": True})
            continue
        trailing_ok = (row.get("bot_trailing_expr_10d") or 0.0) > 0
        tape_ok = row.get("tape_regime") in _tape_favorable_set(prior)
        vol_ok = _vol_side_ok(prior, row.get("spy_vol20"))
        favorable = bool(trailing_ok and tape_ok and vol_ok)
        out.append({"session_date": row["session_date"],
                    "p_profit": 1.0 if favorable else 0.0, "favorable": favorable})
    return out


def full_sample_fit_threshold(rows: list[dict]) -> list[dict]:
    """DELIBERATELY LEAKY variant of score_threshold: uses the WHOLE row set
    (including D and everything after) to derive the tape/vol sign instead of
    `rows[:i]`. Step 4 sanity check ONLY — never the shipped flag. If this
    variant crushes the walk-forward one on held-out expR, the walk-forward
    gate's edge is real; if it doesn't, something upstream is leaking."""
    out = []
    for row in rows:
        trailing_ok = (row.get("bot_trailing_expr_10d") or 0.0) > 0
        tape_ok = row.get("tape_regime") in _tape_favorable_set(rows)
        vol_ok = _vol_side_ok(rows, row.get("spy_vol20"))
        favorable = bool(trailing_ok and tape_ok and vol_ok)
        out.append({"session_date": row["session_date"],
                    "p_profit": 1.0 if favorable else 0.0, "favorable": favorable})
    return out


# --------------------------------------------------------------------------- #
# Step 2 — meta_v1 (sklearn GBM, walk-forward)
# --------------------------------------------------------------------------- #
def _build_X(rows: list[dict], feature_cols=FEATURE_COLS):
    """The feature matrix passed to the classifier — ONLY feature_cols, missing
    values filled 0.0. Never includes label_expr/label_pos
    (test_regime_features_no_label_leak asserts the column set)."""
    import pandas as pd
    return pd.DataFrame(
        [{c: (r.get(c) if r.get(c) is not None else 0.0) for c in feature_cols} for r in rows])


def score_meta(rows: list[dict], p_min: float = P_MIN_DEFAULT,
              min_train: int = MIN_TRAIN_SESSIONS) -> list[dict]:
    """Walk-forward sklearn.GradientBoostingClassifier(max_depth=3, <=8
    features). Fits on `rows[:i]` only (never index >= i — the leakage guard);
    below `min_train` prior labeled sessions, favorable=True (fail-open, don't
    gate on an unfit model)."""
    from sklearn.ensemble import GradientBoostingClassifier

    out = []
    for i, row in enumerate(rows):
        prior = rows[:i]
        train = [r for r in prior if r.get("label_pos") is not None]
        y_train = [bool(r["label_pos"]) for r in train]
        if len(train) < min_train or len(set(y_train)) < 2:
            out.append({"session_date": row["session_date"], "p_profit": None, "favorable": True})
            continue
        clf = GradientBoostingClassifier(max_depth=3, n_estimators=50, random_state=0)
        clf.fit(_build_X(train), y_train)
        p = float(clf.predict_proba(_build_X([row]))[0, list(clf.classes_).index(True)])
        out.append({"session_date": row["session_date"], "p_profit": p, "favorable": p >= p_min})
    return out


# --------------------------------------------------------------------------- #
# score — DB wrapper
# --------------------------------------------------------------------------- #
def _feature_rows(con) -> list[dict]:
    cols = ("session_date", *FEATURE_COLS, *LABEL_COLS)
    q = f"SELECT {', '.join(cols)} FROM regime_features ORDER BY session_date"
    out = []
    for vals in con.execute(q).fetchall():
        row = dict(zip(cols, vals))
        row["session_date"] = _as_date(row["session_date"])
        out.append(row)
    return out


def score(con, model_version: str, p_min: float = P_MIN_DEFAULT) -> int:
    """Step 1/2 `score`: write regime_scores rows for one model_version, walk-
    forward over the finished regime_features table. Idempotent (PK
    (model_version, session_date) + ON CONFLICT DO UPDATE)."""
    apply_schema(con)
    rows = _feature_rows(con)
    if model_version == "threshold_v1":
        scored = score_threshold(rows)
    elif model_version == "meta_v1":
        scored = score_meta(rows, p_min=p_min)
    else:
        raise ValueError(f"unknown model_version: {model_version!r}")

    now = datetime.now(timezone.utc)
    for s in scored:
        con.execute(
            "INSERT INTO regime_scores (model_version, session_date, p_profit, favorable, computed_at) "
            "VALUES (?, ?, ?, ?, ?) ON CONFLICT (model_version, session_date) DO UPDATE SET "
            "p_profit=EXCLUDED.p_profit, favorable=EXCLUDED.favorable, computed_at=EXCLUDED.computed_at",
            [model_version, s["session_date"], s["p_profit"], s["favorable"], now])
    return len(scored)


# --------------------------------------------------------------------------- #
# Step 4 — OOS split + leakage-audit helpers (supervised run wires these up;
# never executed against production here).
# --------------------------------------------------------------------------- #
def oos_split(rows: list[dict], frac: float = 0.3) -> tuple[list[dict], list[dict]]:
    """(head, held_out_tail) — `rows` sorted by session_date ascending; the
    tail is the most-recent `frac` of sessions."""
    n = len(rows)
    k = int(round(n * (1 - frac)))
    return rows[:k], rows[k:]


def sequential_folds(rows: list[dict], k: int = 3) -> list[list[dict]]:
    """`rows` (sorted ascending) split into k sequential chronological folds of
    ~equal size, for the 'expR_delta positive across ~3 sequential folds' OOS
    check."""
    n = len(rows)
    size = max(1, -(-n // k))   # ceil(n/k) so k folds cover all rows
    return [rows[i:i + size] for i in range(0, n, size)][:k]


def session_expr_delta(rows: list[dict], scored: list[dict]) -> float | None:
    """Fire-free session-level expR_delta proxy: mean(label_expr over
    favorable sessions) - mean(label_expr over ALL sessions in `rows`). Usable
    directly on regime_features rows (no experiment_runs needed) for the
    OOS/fold/leakage checks in this step."""
    fav = {s["session_date"] for s in scored if s.get("favorable")}
    fav_labels = [r["label_expr"] for r in rows if r["session_date"] in fav and r.get("label_expr") is not None]
    all_labels = [r["label_expr"] for r in rows if r.get("label_expr") is not None]
    if not fav_labels or not all_labels:
        return None
    return mean(fav_labels) - mean(all_labels)


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def _store():
    from storage.event_store import EventStore
    return EventStore(os.environ.get("WATCHER_EVENT_DB_PATH", "momentum"))


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(prog="python -m research.regime")
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("probe")
    bf = sub.add_parser("build-features")
    bf.add_argument("--since", help="only rebuild sessions >= this date (YYYY-MM-DD)")
    sc = sub.add_parser("score")
    sc.add_argument("--model", required=True, choices=["threshold_v1", "meta_v1"])
    sc.add_argument("--p-min", type=float, default=P_MIN_DEFAULT)
    args = ap.parse_args(argv)

    store = _store()
    try:
        if args.cmd == "probe":
            info = probe(store.con)
            for k, v in info.items():
                print(f"{k}: {v}")
            if info["stop"]:
                print(f"STOP: {info['labelable_sessions']} labelable sessions "
                      f"< {MIN_LABELABLE_SESSIONS} (thin sample)")
                return 1
            return 0
        if args.cmd == "build-features":
            n = build_features(store.con, since=args.since)
            print(f"regime_features: {n} rows")
            return 0
        if args.cmd == "score":
            n = score(store.con, args.model, p_min=args.p_min)
            print(f"regime_scores[{args.model}]: {n} rows")
            return 0
        return 2
    finally:
        store.close()


if __name__ == "__main__":
    raise SystemExit(main())