fithia2/scripts/orb_continuation_labels.py

"""ORB Continuation-Value Labels via Longstaff-Schwartz Backward Induction.

Reads the Phase B per-bar state dataset and computes, for each (trade, bar):

    sell_now_r(t)   = next-bar open R (executable price if you decide at t close)
    continue_r(t)   = cross-fitted ĉ(state_t) — predicted value of *not selling now*
    V_t             = max(sell_now_r(t), continue_r(t))      # LSM optimal value
    optimal_action  = "sell" if sell_now_r(t) >= continue_r(t) else "hold"

Backward induction:

    V_{T_i}        = sell_now_r(T_i)                          # last bar: forced sale
    V_t            = max(sell_now_r(t), ĉ(state_t))           # for t < T_i

ĉ is a single LightGBM regressor trained per backward step on (state_t, V_{t+1})
pairs across all trades with a bar at index t+1, with **trade-level K-fold
cross-fitting** so each trade's ĉ prediction comes from a model not trained on
that trade.

Features intentionally exclude any "future-looking" columns. Trade-level static
context is broadcast to each bar.

Usage:
    python scripts/orb_continuation_labels.py \
        --input tmp/orb_perbar_v49_91_200d/per_bar_states.parquet \
        --out tmp/orb_perbar_v49_91_200d_lsm
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
import lightgbm as lgb
from sklearn.model_selection import KFold


# ---------- Feature configuration ----------

# Per-bar dynamic state features (visible at decision time t).
DYN_FEATURES = [
    "minutes_since_entry",
    "minutes_to_close",
    "minutes_since_open",
    "bar_return_pct",
    "bar_close_loc",
    "current_close_r",
    "current_high_r",
    "current_low_r",
    "mfe_so_far_r",
    "mae_so_far_r",
    "giveback_from_peak_r",
    "bars_since_peak",
    "vwap_dev_pct",
    "vol_vs_first_bar",
    "vol_vs_avg_dvol30d",
    "spy_return_since_entry",
    "spy_return_since_open",
    "qqq_return_since_entry",
    "qqq_return_since_open",
]

# Trade-level static features (broadcast — known at entry, fixed for trade).
STATIC_FEATURES = [
    "atr_at_entry",
    "gap_pct",
    "rvol",
    "morning_gain_pct",
    "entropy_20d",
    "ret_5d",
    "candidate_score",
    "score_rank_pct",
    "sector_confirmation_score",
    "entry_dollar_volume",
    "avg_dollar_vol_30d",
    "premarket_dollar_vol",
    "first_bar_dollar_vol",
    "body_ratio",
    "close_location",
    "gap_zscore_20d",
    "obv_slope_20",
    "obv_slope_5",
    "orb_return",
]

ALL_FEATURES = DYN_FEATURES + STATIC_FEATURES


def lgb_params(seed: int = 0) -> dict:
    return {
        "objective": "regression",
        "metric": "rmse",
        "learning_rate": 0.05,
        "num_leaves": 31,
        "min_data_in_leaf": 30,
        "feature_fraction": 0.9,
        "bagging_fraction": 0.8,
        "bagging_freq": 1,
        "verbose": -1,
        "seed": seed,
    }


def cross_fit_predict(
    X: np.ndarray, y: np.ndarray, trade_ids: np.ndarray, n_splits: int = 5, seed: int = 0
) -> np.ndarray:
    """Trade-level K-fold cross-fit: each trade's prediction comes from a fold not containing it.

    With small n_trades, num_boost_round is kept short and trees shallow.
    """
    unique_trades = np.unique(trade_ids)
    kf = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
    preds = np.zeros(len(y))

    for fold, (train_t_idx, val_t_idx) in enumerate(kf.split(unique_trades)):
        train_trades = set(unique_trades[train_t_idx])
        val_trades = set(unique_trades[val_t_idx])
        train_mask = np.isin(trade_ids, list(train_trades))
        val_mask = np.isin(trade_ids, list(val_trades))

        if train_mask.sum() < 50 or val_mask.sum() == 0:
            preds[val_mask] = y.mean()
            continue

        ds = lgb.Dataset(X[train_mask], label=y[train_mask])
        model = lgb.train(
            lgb_params(seed + fold),
            ds,
            num_boost_round=120,
        )
        preds[val_mask] = model.predict(X[val_mask])

    return preds


def run_lsm(
    per_bar: pd.DataFrame,
    n_splits: int = 5,
    max_iters: int = 8,
    tol: float = 0.02,
    seed: int = 0,
) -> tuple[pd.DataFrame, list[float]]:
    """LSM via global-pooled fitted-value-iteration.

    Pools (state_t, V_{t+1}) pairs across all (trade, bar) rather than fitting one
    model per backward step (which underfits on 146 trades). At each outer iter:
        1. Compute V_next per row via groupby-shift within trade.
        2. Cross-fit a global LightGBM on (state, V_next) with trade-level K-fold.
        3. continue_r ← cross-fit prediction; V ← max(sell_now_r, continue_r).
    Iterate until max |ΔV| < tol.

    Returns (labeled_df, [delta_history]).
    """
    df = per_bar.sort_values(["trade_id", "bar_idx"]).reset_index(drop=True).copy()
    df["sell_now_r"] = df["next_open_r"].astype(float)

    grp = df.groupby("trade_id")
    df["n_bars"] = grp["bar_idx"].transform("size").astype(int)
    df["is_last_bar"] = (df["bar_idx"] == df["n_bars"] - 1).astype(bool)

    # Initial V: assume immediate sale at every bar (a defensible starting policy).
    df["V"] = df["sell_now_r"].copy()
    df["continue_r"] = np.nan

    X_full = df[ALL_FEATURES].astype(float).to_numpy()
    trade_ids_full = df["trade_id"].to_numpy()

    deltas: list[float] = []
    sn = df["sell_now_r"].to_numpy()
    is_last = df["is_last_bar"].to_numpy()

    for it in range(max_iters):
        # V_next within trade (NaN for the last bar of each trade).
        df["V_next"] = df.groupby("trade_id")["V"].shift(-1)
        valid_mask = df["V_next"].notna().to_numpy()
        if valid_mask.sum() < n_splits * 50:
            print(f"  iter {it+1}: not enough samples ({int(valid_mask.sum())}); abort")
            break

        X = X_full[valid_mask]
        y = df.loc[valid_mask, "V_next"].to_numpy()
        tids = trade_ids_full[valid_mask]

        preds = cross_fit_predict(X, y, tids, n_splits=n_splits, seed=seed + it)
        cr = np.full(len(df), np.nan)
        cr[valid_mask] = preds
        df["continue_r"] = cr

        # V_t = max(sell_now_r, continue_r) for non-last bars; sell_now_r at last bar.
        new_V = sn.copy()
        non_last_with_pred = (~is_last) & ~np.isnan(cr)
        new_V[non_last_with_pred] = np.maximum(sn[non_last_with_pred], cr[non_last_with_pred])

        delta = float(np.nanmax(np.abs(new_V - df["V"].to_numpy())))
        deltas.append(delta)
        df["V"] = new_V
        print(f"  iter {it+1}: max |ΔV| = {delta:.4f}")
        if delta < tol:
            print(f"  converged at iter {it+1}")
            break

    df["optimal_action"] = np.where(
        is_last,
        "sell",
        np.where(
            sn >= df["continue_r"].fillna(-np.inf).to_numpy(),
            "sell",
            "hold",
        ),
    )
    df.drop(columns=["V_next"], errors="ignore", inplace=True)
    return df, deltas


def policy_simulate(df: pd.DataFrame) -> pd.DataFrame:
    """For each trade, walk forward bar-by-bar and apply the LSM policy:
    sell at first bar where optimal_action == 'sell' (or last bar). Returns
    one row per trade with policy realized R and stopping bar."""
    rows = []
    for tid, g in df.sort_values(["trade_id", "bar_idx"]).groupby("trade_id"):
        g = g.reset_index(drop=True)
        first_sell = g.index[g["optimal_action"] == "sell"]
        idx = int(first_sell[0]) if len(first_sell) else len(g) - 1
        rows.append(
            {
                "trade_id": tid,
                "ticker": g.iloc[0]["ticker"],
                "date": g.iloc[0]["date"],
                "policy_exit_bar_idx": idx,
                "policy_exit_minutes_since_entry": float(g.iloc[idx]["minutes_since_entry"]),
                "policy_realized_r": float(g.iloc[idx]["sell_now_r"]),
                "baseline_exit_bar_idx": int(g.iloc[0]["baseline_exit_bar_idx"]),
                "baseline_realized_r": float(g.iloc[0]["baseline_realized_r"]),
                "n_bars": int(g.iloc[0]["n_bars"]),
                "mfe_r_full_path": float(g["mfe_so_far_r"].iloc[-1]),
                "oracle_close_r_full_path": float(g["current_close_r"].max()),
            }
        )
    return pd.DataFrame(rows)


def build_summary(per_bar: pd.DataFrame, policy_df: pd.DataFrame) -> str:
    base_total = float(policy_df["baseline_realized_r"].sum())
    base_mean = float(policy_df["baseline_realized_r"].mean())
    pol_total = float(policy_df["policy_realized_r"].sum())
    pol_mean = float(policy_df["policy_realized_r"].mean())
    delta = policy_df["policy_realized_r"] - policy_df["baseline_realized_r"]

    # Feature-importance: train one supervisor model on (state_t, V) for diagnostics.
    valid = per_bar.dropna(subset=["continue_r"]).copy()
    X = valid[ALL_FEATURES].astype(float).to_numpy()
    y = valid["V"].to_numpy()
    ds = lgb.Dataset(X, label=y, feature_name=ALL_FEATURES)
    diag_model = lgb.train(lgb_params(seed=42), ds, num_boost_round=200)
    importances = sorted(
        zip(ALL_FEATURES, diag_model.feature_importance(importance_type="gain")),
        key=lambda x: -x[1],
    )

    n_trades = policy_df["trade_id"].nunique()
    sell_rows = (per_bar["optimal_action"] == "sell").sum()
    hold_rows = (per_bar["optimal_action"] == "hold").sum()

    lines = [
        "# Phase C — LSM Continuation Labels",
        "",
        f"Trades: **{n_trades}** | per-bar rows: **{len(per_bar):,}** "
        f"(sell labels: {sell_rows:,}, hold labels: {hold_rows:,})",
        "",
        "## In-sample LSM policy vs V49.91 baseline",
        "",
        f"- baseline total_R: **{base_total:.2f}** (mean {base_mean:.3f})",
        f"- LSM in-sample total_R: **{pol_total:.2f}** (mean {pol_mean:.3f})",
        f"- Δtotal_R: **{pol_total - base_total:+.2f}** "
        f"(mean Δ per trade {(delta.mean()):+.3f})",
        f"- Trades where LSM beats baseline: "
        f"**{int((delta > 0).sum())}/{len(delta)} ({(delta > 0).mean():.1%})**",
        f"- Trades where LSM ties: **{int((delta == 0).sum())}**",
        "",
        "_⚠ In-sample numbers — Phase D walk-forward is the real test._",
        "",
        "## Policy hold-time vs baseline (bars from entry)",
        "",
        "| stat | baseline | LSM policy |",
        "|---|---|---|",
        f"| mean | {policy_df['baseline_exit_bar_idx'].mean():.1f} | "
        f"{policy_df['policy_exit_bar_idx'].mean():.1f} |",
        f"| median | {policy_df['baseline_exit_bar_idx'].median():.0f} | "
        f"{policy_df['policy_exit_bar_idx'].median():.0f} |",
        f"| p25 | {policy_df['baseline_exit_bar_idx'].quantile(0.25):.0f} | "
        f"{policy_df['policy_exit_bar_idx'].quantile(0.25):.0f} |",
        f"| p75 | {policy_df['baseline_exit_bar_idx'].quantile(0.75):.0f} | "
        f"{policy_df['policy_exit_bar_idx'].quantile(0.75):.0f} |",
        "",
        "## Feature importance (gain, supervisor diagnostic on state→V)",
        "",
        "| rank | feature | gain |",
        "|---|---|---|",
    ]
    for i, (f, g) in enumerate(importances[:20], 1):
        lines.append(f"| {i} | `{f}` | {g:.0f} |")
    return "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--n-splits", type=int, default=5)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    per_bar = pd.read_parquet(args.input)
    print(f"Loaded {len(per_bar):,} bar rows from {per_bar['trade_id'].nunique()} trades")

    print("Running LSM via global-pooled fitted value iteration...")
    labeled, deltas = run_lsm(per_bar, n_splits=args.n_splits, seed=args.seed)
    out_parquet = out_dir / "per_bar_lsm.parquet"
    labeled.to_parquet(out_parquet, index=False)
    print(f"Wrote {out_parquet} (delta history: {[round(d,4) for d in deltas]})")

    policy_df = policy_simulate(labeled)
    pol_parquet = out_dir / "policy_per_trade.parquet"
    policy_df.to_parquet(pol_parquet, index=False)
    print(f"Wrote {pol_parquet}")

    md = build_summary(labeled, policy_df)
    (out_dir / "summary.md").write_text(md)
    print(f"Wrote {out_dir / 'summary.md'}")


if __name__ == "__main__":
    main()