fithia2/scripts/orb_phase_d_clean.py

"""Phase D (clean) — Walk-forward with no LSM label leakage.

Train: clean LSM labels from train-only trades.
Test: raw per-bar states (no LSM labels needed — model predicts continue_r directly).

Also runs a "no-stop-override" variant: trains only on normal-exit trades (baseline
bar > 5) to isolate the pure exit-timing signal from stop-override signal.

Outputs:
    model_clean.txt              — model trained on clean train-only LSM labels
    model_no_stop.txt            — model trained without stop-override trades
    test_clean.parquet           — test results with clean model
    test_no_stop.parquet         — test results with no-stop model
    summary.md

Usage:
    python scripts/orb_phase_d_clean.py \
        --train-lsm tmp/orb_phase_d_clean/train_lsm/per_bar_lsm.parquet \
        --test-states tmp/orb_phase_d_clean/test_per_bar_states.parquet \
        --out tmp/orb_phase_d_clean
"""
from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
import lightgbm as lgb


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",
]

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.03,
        "num_leaves": 31,
        "min_data_in_leaf": 20,
        "feature_fraction": 0.8,
        "bagging_fraction": 0.8,
        "bagging_freq": 5,
        "verbose": -1,
        "seed": seed,
    }


def train_model(train_df: pd.DataFrame, exclude_early_stop: bool = False, seed: int = 0) -> lgb.Booster:
    rows = train_df[train_df["is_last_bar"] == False].dropna(subset=["continue_r"]).copy()
    if exclude_early_stop:
        # Exclude trades where baseline exited early (bar <= 5) — isolate exit-timing signal
        early_stop_ids = train_df.groupby("trade_id").first()
        early_stop_ids = early_stop_ids[early_stop_ids["baseline_exit_bar_idx"] <= 5].index
        rows = rows[~rows["trade_id"].isin(early_stop_ids)]
    label = "no-stop" if exclude_early_stop else "clean"
    print(f"  [{label}] Training on {len(rows):,} rows from {rows['trade_id'].nunique()} trades")
    X = rows[ALL_FEATURES].astype(float).values
    y = rows["continue_r"].astype(float).values
    ds = lgb.Dataset(X, label=y, feature_name=ALL_FEATURES)
    return lgb.train(lgb_params(seed), ds, num_boost_round=300, callbacks=[lgb.log_evaluation(100)])


def prepare_test(test_states: pd.DataFrame) -> pd.DataFrame:
    """Add derived columns needed for policy simulation."""
    df = test_states.sort_values(["trade_id", "bar_idx"]).copy()
    df["sell_now_r"] = df["next_open_r"].astype(float)
    n_bars = df.groupby("trade_id")["bar_idx"].transform("size").astype(int)
    df["n_bars"] = n_bars
    df["is_last_bar"] = (df["bar_idx"] == n_bars - 1).astype(bool)
    return df


def apply_policy(model: lgb.Booster, df: pd.DataFrame, hard_stop_r: float = -2.0) -> pd.DataFrame:
    """Apply model exit policy with an optional hard stop.

    hard_stop_r: exit immediately if sell_now_r falls below this threshold (e.g., -2.0R).
    Set to -inf to disable hard stop.
    """
    rows_out = []
    for tid, g in df.sort_values(["trade_id", "bar_idx"]).groupby("trade_id"):
        g = g.reset_index(drop=True)
        X = g[ALL_FEATURES].astype(float).values
        pred_cr = model.predict(X)
        n = len(g)

        exit_idx = n - 1
        for i in range(n):
            sn = float(g.iloc[i]["sell_now_r"])
            is_last = bool(g.iloc[i]["is_last_bar"])
            # hard stop fires first
            if sn <= hard_stop_r:
                exit_idx = i
                break
            if is_last or sn >= float(pred_cr[i]):
                exit_idx = i
                break

        rows_out.append({
            "trade_id": tid,
            "ticker": g.iloc[0]["ticker"],
            "date": g.iloc[0]["date"],
            "model_exit_bar_idx": exit_idx,
            "model_exit_minutes": float(g.iloc[exit_idx]["minutes_since_entry"]),
            "model_realized_r": float(g.iloc[exit_idx]["sell_now_r"]),
            "baseline_exit_bar_idx": float(g.iloc[0]["baseline_exit_bar_idx"]),
            "baseline_realized_r": float(g.iloc[0]["baseline_realized_r"]),
            "n_bars": n,
            "mfe_r": float(g["mfe_so_far_r"].iloc[-1]),
        })
    return pd.DataFrame(rows_out)


def sharpe(r: pd.Series) -> float:
    return float(r.mean() / r.std() * np.sqrt(252)) if r.std() > 0 else float("nan")


def max_dd(r: pd.Series) -> float:
    cum = r.cumsum()
    return float((cum - cum.cummax()).min())


def section(df: pd.DataFrame, label: str) -> list[str]:
    lines = [f"## {label} — {len(df)} test trades", ""]
    base = df["baseline_realized_r"]
    mod = df["model_realized_r"]
    delta = mod - base
    lines += [
        "| metric | baseline | model | delta |",
        "|---|---|---|---|",
        f"| total_R | {base.sum():.2f} | {mod.sum():.2f} | {delta.sum():+.2f} |",
        f"| mean_R | {base.mean():.3f} | {mod.mean():.3f} | {delta.mean():+.3f} |",
        f"| Sharpe | {sharpe(base):.2f} | {sharpe(mod):.2f} | — |",
        f"| max_DD | {max_dd(base):.3f} | {max_dd(mod):.3f} | — |",
        f"| beats baseline | — | {int((delta>0).sum())}/{len(df)} ({(delta>0).mean():.1%}) | — |",
        "",
        "**Decomposition by baseline exit type:**",
        "",
        "| category | n | base Σ R | model Σ R | delta |",
        "|---|---|---|---|---|",
    ]
    for bar_thresh, cat in [(5, "Early-stop (bar≤5)"), (9999, "Normal exit (bar>5)")]:
        if bar_thresh == 5:
            sub = df[df["baseline_exit_bar_idx"] <= 5]
        else:
            sub = df[df["baseline_exit_bar_idx"] > 5]
        if len(sub):
            d = sub["model_realized_r"] - sub["baseline_realized_r"]
            lines.append(f"| {cat} | {len(sub)} | {sub['baseline_realized_r'].sum():.2f} | "
                         f"{sub['model_realized_r'].sum():.2f} | {d.sum():+.2f} |")
    lines.append("")
    return lines


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--train-lsm", required=True)
    ap.add_argument("--test-states", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    out = Path(args.out)

    train_lsm = pd.read_parquet(args.train_lsm)
    test_raw = pd.read_parquet(args.test_states)
    test_df = prepare_test(test_raw)
    print(f"Train: {train_lsm['trade_id'].nunique()} trades ({len(train_lsm):,} rows)")
    print(f"Test:  {test_df['trade_id'].nunique()} trades ({len(test_df):,} rows)")

    # --- Model 1: clean (all train trades, no leakage), no hard stop ---
    print("\nTraining model 1 (clean)...")
    m_clean = train_model(train_lsm, exclude_early_stop=False, seed=args.seed)
    m_clean.save_model(str(out / "model_clean.txt"))
    res_clean = apply_policy(m_clean, test_df, hard_stop_r=float("-inf"))
    res_clean.to_parquet(out / "test_clean.parquet", index=False)
    print(f"Clean model (no hard stop): total_R={res_clean['model_realized_r'].sum():.2f} "
          f"(baseline {res_clean['baseline_realized_r'].sum():.2f})")

    # --- Model 1b: clean + hard stop at -2R ---
    res_clean_stop = apply_policy(m_clean, test_df, hard_stop_r=-2.0)
    res_clean_stop.to_parquet(out / "test_clean_hardstop.parquet", index=False)
    print(f"Clean model (+2R hard stop):  total_R={res_clean_stop['model_realized_r'].sum():.2f}")

    # --- Model 2: no-stop-override (only normal-exit train trades) ---
    print("\nTraining model 2 (no stop-override trades)...")
    m_nostop = train_model(train_lsm, exclude_early_stop=True, seed=args.seed)
    m_nostop.save_model(str(out / "model_no_stop.txt"))
    res_nostop = apply_policy(m_nostop, test_df, hard_stop_r=float("-inf"))
    res_nostop.to_parquet(out / "test_no_stop.parquet", index=False)
    print(f"No-stop model (no hard stop): total_R={res_nostop['model_realized_r'].sum():.2f} "
          f"(baseline {res_nostop['baseline_realized_r'].sum():.2f})")

    # --- Model 2b: no-stop + hard stop at -2R ---
    res_nostop_stop = apply_policy(m_nostop, test_df, hard_stop_r=-2.0)
    res_nostop_stop.to_parquet(out / "test_no_stop_hardstop.parquet", index=False)
    print(f"No-stop model (+2R hard stop): total_R={res_nostop_stop['model_realized_r'].sum():.2f}")

    # --- Feature importances ---
    imp_clean = sorted(
        zip(ALL_FEATURES, m_clean.feature_importance("gain")), key=lambda x: -x[1]
    )
    imp_nostop = sorted(
        zip(ALL_FEATURES, m_nostop.feature_importance("gain")), key=lambda x: -x[1]
    )

    # --- Summary ---
    lines = [
        "# Phase D (Clean) — Walk-Forward, No LSM Label Leakage",
        "",
        f"Train: {train_lsm['trade_id'].nunique()} trades (≤2026-02-01, LSM on train-only)",
        f"Test: {test_df['trade_id'].nunique()} trades (>2026-02-01, no LSM labels used)",
        "",
        "_Reference: original Phase D (with leakage) test Δ = +4.74R — was inflated by LSM leakage_",
        "",
    ]
    lines += section(res_clean, "Model 1a — Clean, no hard stop")
    lines += section(res_clean_stop, "Model 1b — Clean + hard stop at −2R")
    lines += section(res_nostop, "Model 2a — No-stop-override training, no hard stop")
    lines += section(res_nostop_stop, "Model 2b — No-stop-override training + hard stop at −2R")

    lines += [
        "## Feature importance comparison (top 15, gain)",
        "",
        "| rank | clean model | no-stop model |",
        "|---|---|---|",
    ]
    for i, ((f1, g1), (f2, g2)) in enumerate(zip(imp_clean[:15], imp_nostop[:15]), 1):
        lines.append(f"| {i} | `{f1}` ({g1:.0f}) | `{f2}` ({g2:.0f}) |")

    md = "\n".join(lines)
    (out / "summary_clean.md").write_text(md)
    print(f"\nWrote {out / 'summary_clean.md'}")
    print("\n" + md)


if __name__ == "__main__":
    main()