"""Phase D — Walk-forward evaluation of the LSM continuation model.

Trains a LightGBM regressor on LSM-labeled per-bar data (train set: trades ≤ CUTOFF)
to predict the continuation value, then applies the trained model to the test set
(trades > CUTOFF) to make exit decisions.

Decision rule at each bar: sell when sell_now_r >= predicted_continue_r.

Outputs:
    model.txt                   — trained LightGBM model
    test_per_trade.parquet      — one row per test trade: model vs baseline
    train_per_trade.parquet     — one row per train trade: model vs baseline (in-sample)
    summary.md                  — comparison table + feature importance

Usage:
    python scripts/orb_phase_d_walkforward.py \
        --input tmp/orb_perbar_v49_91_200d_lsm/per_bar_lsm.parquet \
        --out tmp/orb_phase_d \
        --cutoff 2026-02-01
"""
from __future__ import annotations

import argparse
from pathlib import Path

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


# Same feature set as Phase C (V49 entries have all, including candidate_score).
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, seed: int = 0) -> lgb.Booster:
    """Train LightGBM regressor on (state, continue_r) pairs from non-last bars."""
    rows = train_df[train_df["is_last_bar"] == False].dropna(subset=["continue_r"])
    print(f"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)
    model = lgb.train(
        lgb_params(seed),
        ds,
        num_boost_round=300,
        callbacks=[lgb.log_evaluation(50)],
    )
    return model


def apply_policy(model: lgb.Booster, df: pd.DataFrame) -> pd.DataFrame:
    """For each trade, walk forward and sell at first bar where sell_now_r >= predicted_continue_r.

    Returns one row per trade with model exit vs baseline.
    """
    rows_out = []
    for tid, g in df.sort_values(["trade_id", "bar_idx"]).groupby("trade_id"):
        g = g.reset_index(drop=True)
        n = len(g)

        X = g[ALL_FEATURES].astype(float).values
        pred_cr = model.predict(X)

        exit_idx = n - 1  # default: last bar
        for i in range(n):
            sn = float(g.iloc[i]["sell_now_r"])
            cr = float(pred_cr[i])
            is_last = bool(g.iloc[i]["is_last_bar"])
            if is_last or sn >= cr:
                exit_idx = i
                break

        baseline_r = g.iloc[0]["baseline_realized_r"]
        baseline_bar = g.iloc[0]["baseline_exit_bar_idx"]

        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_since_entry": float(g.iloc[exit_idx]["minutes_since_entry"]),
            "model_realized_r": float(g.iloc[exit_idx]["sell_now_r"]),
            "baseline_exit_bar_idx": float(baseline_bar),
            "baseline_realized_r": float(baseline_r),
            "n_bars": n,
            "lsm_optimal_bar": int(g.index[g["optimal_action"] == "sell"][0])
            if (g["optimal_action"] == "sell").any()
            else n - 1,
            "lsm_realized_r": float(
                g.iloc[
                    int(g.index[g["optimal_action"] == "sell"][0])
                    if (g["optimal_action"] == "sell").any()
                    else n - 1
                ]["sell_now_r"]
            ),
            "mfe_r": float(g["mfe_so_far_r"].iloc[-1]),
            "oracle_r": float(g["current_close_r"].max()),
        })

    return pd.DataFrame(rows_out)


def sharpe_ratio(returns: pd.Series) -> float:
    if len(returns) < 2 or returns.std() == 0:
        return float("nan")
    return float(returns.mean() / returns.std() * np.sqrt(252))


def max_drawdown(returns: pd.Series) -> float:
    cum = returns.cumsum()
    running_max = cum.cummax()
    dd = cum - running_max
    return float(dd.min())


def build_summary(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    model: lgb.Booster,
) -> str:
    lines = ["# Phase D — Walk-Forward Evaluation", ""]

    for label, df in [("Train (in-sample)", train_df), ("Test (walk-forward)", test_df)]:
        n = len(df)
        valid = df.dropna(subset=["baseline_realized_r"])

        model_r = valid["model_realized_r"]
        base_r = valid["baseline_realized_r"]
        delta = model_r - base_r

        lines += [
            f"## {label} — {n} trades",
            "",
            "| metric | baseline | model | delta |",
            "|---|---|---|---|",
            f"| total_R | {base_r.sum():.2f} | {model_r.sum():.2f} | {delta.sum():+.2f} |",
            f"| mean_R | {base_r.mean():.3f} | {model_r.mean():.3f} | {delta.mean():+.3f} |",
            f"| Sharpe (trade-level) | {sharpe_ratio(base_r):.2f} | {sharpe_ratio(model_r):.2f} | — |",
            f"| max_DD | {max_drawdown(base_r):.3f} | {max_drawdown(model_r):.3f} | — |",
            f"| model beats baseline | — | {int((delta>0).sum())}/{n} ({(delta>0).mean():.1%}) | — |",
            "",
            "| hold-time stat | baseline | model |",
            "|---|---|---|",
            f"| mean bars | {valid['baseline_exit_bar_idx'].mean():.1f} | {valid['model_exit_bar_idx'].mean():.1f} |",
            f"| median bars | {valid['baseline_exit_bar_idx'].median():.0f} | {valid['model_exit_bar_idx'].median():.0f} |",
            f"| p25 bars | {valid['baseline_exit_bar_idx'].quantile(0.25):.0f} | {valid['model_exit_bar_idx'].quantile(0.25):.0f} |",
            f"| p75 bars | {valid['baseline_exit_bar_idx'].quantile(0.75):.0f} | {valid['model_exit_bar_idx'].quantile(0.75):.0f} |",
            "",
        ]

    # Feature importance
    importances = sorted(
        zip(ALL_FEATURES, model.feature_importance(importance_type="gain")),
        key=lambda x: -x[1],
    )
    lines += [
        "## Feature importance (gain)",
        "",
        "| 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("--cutoff", default="2026-02-01", help="Train/test split date (inclusive train)")
    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)

    df = pd.read_parquet(args.input)
    print(f"Loaded {len(df):,} rows from {df['trade_id'].nunique()} trades "
          f"({df['date'].min()} → {df['date'].max()})")

    # Time split by trade date
    trade_first_date = df.groupby("trade_id")["date"].first()
    train_ids = trade_first_date[trade_first_date <= args.cutoff].index
    test_ids = trade_first_date[trade_first_date > args.cutoff].index

    train_df_rows = df[df["trade_id"].isin(train_ids)]
    test_df_rows = df[df["trade_id"].isin(test_ids)]
    print(f"Train: {len(train_ids)} trades ({len(train_df_rows):,} rows)")
    print(f"Test:  {len(test_ids)} trades ({len(test_df_rows):,} rows)")

    print("\nTraining model...")
    model = train_model(train_df_rows, seed=args.seed)
    model.save_model(str(out_dir / "model.txt"))
    print(f"Saved model to {out_dir / 'model.txt'}")

    print("\nApplying policy to train set (in-sample)...")
    train_results = apply_policy(model, train_df_rows)
    train_results.to_parquet(out_dir / "train_per_trade.parquet", index=False)

    print("Applying policy to test set (walk-forward)...")
    test_results = apply_policy(model, test_df_rows)
    test_results.to_parquet(out_dir / "test_per_trade.parquet", index=False)

    md = build_summary(train_results, test_results, model)
    (out_dir / "summary.md").write_text(md)
    print(f"\nWrote {out_dir / 'summary.md'}")
    print(md)


if __name__ == "__main__":
    main()