#!/usr/bin/env python3
"""
Meta-model: learn an empirical "will this trade win?" probability from
backtest output, replacing the engine's heuristic A+/A/B grades with a
calibrated win-probability score.

Inputs:
    backtest_results/trades.csv  (or --trades <path>)

Outputs (under backtest_results/meta/):
    feature_importances.txt
    classification_report.txt
    calibration.png              — predicted prob vs actual win rate
    decile_returns.png           — avg R per predicted-prob decile
    equity_compare.png           — engine grades vs meta-top-N
    summary.txt                  — headline numbers + recommended threshold

Train/test split is **time-based** (first 70% train, last 30% test) — random
splits leak future regime information into training and overstate test
performance on financial time series.

Usage:
    python meta_model.py
    python meta_model.py --trades backtest_results_h1/trades.csv --out meta_h1
    python meta_model.py --top-decile 0.2   # take top 20% predictions
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path
from typing import Tuple

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


def engineer_features(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.Series]:
    """
    Build features available at trade entry time. Strict no-leakage:
    nothing from exit_time, exit_price, outcome, bars_held, r_multiple.
    """
    df = df.copy()
    df["entry_time"] = pd.to_datetime(df["entry_time"])

    # Time-of-day / day-of-week
    df["hour"] = df["entry_time"].dt.hour
    df["dow"]  = df["entry_time"].dt.dayofweek
    df["is_buy"] = (df["direction"] == "BUY").astype(int)

    # MTF flag
    df["mtf"] = df["mtf_aligned"].astype(int)

    # Score / confidence — keep raw and signed/abs versions; model picks
    df["abs_score"] = df["score"].abs()
    df["score_sign"] = np.sign(df["score"])

    # Risk geometry
    df["rr"] = df["reward_pips"] / df["risk_pips"].replace(0, np.nan)
    df["rr"] = df["rr"].fillna(2.0)

    # Regime one-hot
    regime_dummies = pd.get_dummies(df["regime"], prefix="regime").astype(int)
    df = pd.concat([df, regime_dummies], axis=1)

    # Pair one-hot (model can learn pair-specific behavior)
    pair_dummies = pd.get_dummies(df["pair"], prefix="pair").astype(int)
    df = pd.concat([df, pair_dummies], axis=1)

    # Grade one-hot — included so the meta-model can re-weight rather
    # than blindly disregard the engine's view
    grade_dummies = pd.get_dummies(df["grade"], prefix="grade").astype(int)
    df = pd.concat([df, grade_dummies], axis=1)

    feature_cols = (
        ["score", "abs_score", "score_sign", "confidence", "rr",
         "is_buy", "mtf", "hour", "dow", "risk_pips", "reward_pips"]
        + list(regime_dummies.columns)
        + list(pair_dummies.columns)
        + list(grade_dummies.columns)
    )

    X = df[feature_cols].astype(float)
    y = (df["r_multiple"] > 0).astype(int)
    return X, y


def time_based_split(
    X: pd.DataFrame, y: pd.Series, df: pd.DataFrame, train_frac: float = 0.7,
):
    """First 70% by entry_time = train, last 30% = test. Standard for time series."""
    order = df["entry_time"].argsort().values
    cutoff = int(len(order) * train_frac)
    train_idx = order[:cutoff]
    test_idx  = order[cutoff:]
    return (
        X.iloc[train_idx], X.iloc[test_idx],
        y.iloc[train_idx], y.iloc[test_idx],
        df.iloc[train_idx], df.iloc[test_idx],
    )


def evaluate_decile_strategy(
    df_test: pd.DataFrame, p_win: np.ndarray, top_frac: float = 0.2,
) -> dict:
    """
    Take only the top X% by predicted win probability. Compare to taking
    every trade in the test set (the 'engine baseline').
    """
    n = len(df_test)
    cutoff_n = max(1, int(n * top_frac))

    df_eval = df_test.copy()
    df_eval["p_win"] = p_win
    df_eval = df_eval.sort_values("p_win", ascending=False)

    top = df_eval.head(cutoff_n)
    base = df_eval

    return {
        "top_frac":    top_frac,
        "top_n":       len(top),
        "top_wr":      (top["r_multiple"] > 0).mean() * 100,
        "top_avg_r":   top["r_multiple"].mean(),
        "top_total_r": top["r_multiple"].sum(),
        "base_n":      len(base),
        "base_wr":     (base["r_multiple"] > 0).mean() * 100,
        "base_avg_r":  base["r_multiple"].mean(),
        "base_total_r": base["r_multiple"].sum(),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Meta-model on backtest trades.")
    parser.add_argument("--trades", default="backtest_results/trades.csv",
                        help="Path to trades.csv (default backtest_results/trades.csv)")
    parser.add_argument("--out", default="backtest_results/meta",
                        help="Output directory")
    parser.add_argument("--train-frac", type=float, default=0.7)
    parser.add_argument("--top-decile", type=float, default=0.2,
                        help="Top fraction of predictions to 'trade' (default 0.2)")
    parser.add_argument("--min-trades", type=int, default=300,
                        help="Refuse to train below this many trades")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(message)s")

    try:
        from sklearn.linear_model import LogisticRegression
        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.preprocessing import StandardScaler
        from sklearn.pipeline import Pipeline
        from sklearn.metrics import (
            classification_report, roc_auc_score, brier_score_loss,
        )
    except ImportError:
        print("❌ scikit-learn not installed.  pip install scikit-learn")
        return 1

    trades_path = Path(args.trades)
    if not trades_path.exists():
        print(f"❌ {trades_path} not found. Run backtest.py first.")
        return 1

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

    df = pd.read_csv(trades_path)
    df["entry_time"] = pd.to_datetime(df["entry_time"])
    if len(df) < args.min_trades:
        print(f"❌ Only {len(df)} trades — need ≥{args.min_trades} for training.")
        print("   Seed more pairs / longer history, or lower --min-trades.")
        return 1

    print(f"📊 Loaded {len(df):,} trades from {trades_path}")
    print(f"   Pairs: {sorted(df['pair'].unique())}")
    print(f"   Date range: {df['entry_time'].min()} → {df['entry_time'].max()}")
    print(f"   Baseline win rate: {(df['r_multiple'] > 0).mean()*100:.1f}%")
    print(f"   Baseline avg R:   {df['r_multiple'].mean():+.3f}")
    print()

    X, y = engineer_features(df)
    print(f"🔧 Features: {list(X.columns)}")
    print(f"   Shape: {X.shape}, target balance: {y.mean()*100:.1f}% wins\n")

    X_tr, X_te, y_tr, y_te, df_tr, df_te = time_based_split(X, y, df, args.train_frac)
    print(f"📚 Train: {len(X_tr):,} trades  ({df_tr['entry_time'].min().date()} → {df_tr['entry_time'].max().date()})")
    print(f"   Test:  {len(X_te):,} trades  ({df_te['entry_time'].min().date()} → {df_te['entry_time'].max().date()})\n")

    # ── Model 1: Logistic regression (interpretable baseline) ─────────
    lr = Pipeline([
        ("scale", StandardScaler()),
        ("lr", LogisticRegression(max_iter=2000, C=1.0)),
    ])
    lr.fit(X_tr, y_tr)
    p_lr = lr.predict_proba(X_te)[:, 1]

    # ── Model 2: Gradient boosting (non-linear, captures interactions) ─
    gb = GradientBoostingClassifier(
        n_estimators=200, max_depth=3, learning_rate=0.05,
        subsample=0.8, random_state=42,
    )
    gb.fit(X_tr, y_tr)
    p_gb = gb.predict_proba(X_te)[:, 1]

    # ── Metrics ───────────────────────────────────────────────────────
    print("══ Logistic Regression ══")
    print(f"  AUC:    {roc_auc_score(y_te, p_lr):.3f}")
    print(f"  Brier:  {brier_score_loss(y_te, p_lr):.3f}  (lower=better, 0.25=random for 50% prior)")
    print()
    print("══ Gradient Boosting ══")
    print(f"  AUC:    {roc_auc_score(y_te, p_gb):.3f}")
    print(f"  Brier:  {brier_score_loss(y_te, p_gb):.3f}")
    print()

    # Pick the better-AUC model for downstream evaluation
    p_best = p_gb if roc_auc_score(y_te, p_gb) >= roc_auc_score(y_te, p_lr) else p_lr
    best_name = "GradientBoosting" if p_best is p_gb else "LogisticRegression"
    print(f"➡️  Selecting {best_name} for top-decile evaluation\n")

    # ── Decile-strategy comparison ─────────────────────────────────────
    res = evaluate_decile_strategy(df_te, p_best, args.top_decile)
    print(f"══ {int(args.top_decile*100)}% top-pick strategy vs trade-everything baseline (test set) ══")
    print(f"  Trade everything:  {res['base_n']:>5} trades  WR {res['base_wr']:5.1f}%  avgR {res['base_avg_r']:+.3f}  totalR {res['base_total_r']:+8.2f}")
    print(f"  Top {int(args.top_decile*100)}% by p_win:    {res['top_n']:>5} trades  WR {res['top_wr']:5.1f}%  avgR {res['top_avg_r']:+.3f}  totalR {res['top_total_r']:+8.2f}")
    print()

    # ── Per-decile breakdown ──────────────────────────────────────────
    df_eval = df_te.copy()
    df_eval["p_win"] = p_best
    df_eval["decile"] = pd.qcut(df_eval["p_win"], 10, labels=False, duplicates="drop")
    decile_summary = df_eval.groupby("decile").agg(
        n=("r_multiple", "count"),
        avg_p_win=("p_win", "mean"),
        actual_wr=("r_multiple", lambda s: (s > 0).mean()),
        avg_r=("r_multiple", "mean"),
        total_r=("r_multiple", "sum"),
    ).round(3)
    decile_summary.index = decile_summary.index.astype(int) + 1  # 1..10
    print("══ Per-decile breakdown (1=lowest p_win, 10=highest) ══")
    print(decile_summary.to_string())
    print()

    # ── Feature importances ────────────────────────────────────────────
    print("══ Top features by importance (from gradient boosting) ══")
    importances = pd.Series(gb.feature_importances_, index=X.columns)
    top_features = importances.sort_values(ascending=False).head(15)
    for name, imp in top_features.items():
        print(f"  {name:30s}  {imp:.4f}")
    print()

    # ── Save full reports ──────────────────────────────────────────────
    (out_dir / "classification_report.txt").write_text(
        f"Logistic Regression\n{'-'*40}\n"
        f"AUC:   {roc_auc_score(y_te, p_lr):.3f}\n"
        f"Brier: {brier_score_loss(y_te, p_lr):.3f}\n\n"
        f"{classification_report(y_te, (p_lr > 0.5).astype(int))}\n"
        f"\nGradient Boosting\n{'-'*40}\n"
        f"AUC:   {roc_auc_score(y_te, p_gb):.3f}\n"
        f"Brier: {brier_score_loss(y_te, p_gb):.3f}\n\n"
        f"{classification_report(y_te, (p_gb > 0.5).astype(int))}\n"
    )
    (out_dir / "feature_importances.txt").write_text(
        importances.sort_values(ascending=False).to_string()
    )
    (out_dir / "decile_breakdown.csv").write_text(decile_summary.to_csv())
    (out_dir / "summary.txt").write_text(
        f"Meta-model on {len(df):,} trades  "
        f"({df['entry_time'].min()} → {df['entry_time'].max()})\n\n"
        f"Baseline (trade everything):\n"
        f"  win rate: {(df['r_multiple'] > 0).mean()*100:.1f}%\n"
        f"  avg R:    {df['r_multiple'].mean():+.3f}\n"
        f"  total R:  {df['r_multiple'].sum():+.2f}\n\n"
        f"Test set ({len(X_te):,} trades, last 30% by date):\n"
        f"  Trade everything:  WR {res['base_wr']:.1f}%  avgR {res['base_avg_r']:+.3f}  totalR {res['base_total_r']:+.2f}\n"
        f"  Top {int(args.top_decile*100)}% by p_win:  WR {res['top_wr']:.1f}%  avgR {res['top_avg_r']:+.3f}  totalR {res['top_total_r']:+.2f}\n\n"
        f"Best model: {best_name}\n"
        f"  AUC:   {roc_auc_score(y_te, p_best):.3f}  (0.5 = random, >0.55 = real signal)\n"
        f"  Brier: {brier_score_loss(y_te, p_best):.3f}\n"
    )

    # ── Plots ─────────────────────────────────────────────────────────
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    # Calibration curve
    bins = pd.cut(p_best, bins=np.linspace(0, 1, 11))
    cal = pd.DataFrame({"p_win": p_best, "y": y_te.values, "bin": bins})
    grouped = cal.groupby("bin", observed=True).agg(
        mean_pred=("p_win", "mean"),
        actual=("y", "mean"),
        n=("y", "count"),
    )
    fig, ax = plt.subplots(figsize=(6, 5))
    ax.plot([0, 1], [0, 1], "--", color="gray", label="Perfect calibration")
    ax.plot(grouped["mean_pred"], grouped["actual"], "o-", linewidth=2, markersize=8)
    for _, row in grouped.iterrows():
        ax.annotate(f"n={int(row['n'])}",
                    (row["mean_pred"], row["actual"]),
                    textcoords="offset points", xytext=(8, -3), fontsize=8)
    ax.set_xlabel("Mean predicted win probability")
    ax.set_ylabel("Actual win rate")
    ax.set_title(f"Calibration ({best_name}, test set)")
    ax.set_xlim(0, 1); ax.set_ylim(0, 1)
    ax.legend(); ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / "calibration.png", dpi=120)
    plt.close(fig)

    # Decile returns
    fig, ax = plt.subplots(figsize=(8, 5))
    colors = ["#dc2626" if r < 0 else "#16a34a" for r in decile_summary["avg_r"]]
    ax.bar(decile_summary.index.astype(str), decile_summary["avg_r"], color=colors)
    ax.axhline(0, color="black", linewidth=0.5)
    ax.set_xlabel("Predicted-probability decile (1=lowest, 10=highest)")
    ax.set_ylabel("Average R per trade")
    ax.set_title(f"Trade outcome by predicted-probability decile ({best_name})")
    for i, (idx, row) in enumerate(decile_summary.iterrows()):
        ax.text(i, row["avg_r"] + (0.02 if row["avg_r"] >= 0 else -0.05),
                f"n={int(row['n'])}\nWR={row['actual_wr']*100:.0f}%",
                ha="center", fontsize=8)
    ax.grid(alpha=0.3, axis="y")
    fig.tight_layout()
    fig.savefig(out_dir / "decile_returns.png", dpi=120)
    plt.close(fig)

    # Equity comparison: trade-everything vs top-decile (test set, time-ordered)
    df_compare = df_te.copy().sort_values("entry_time").reset_index(drop=True)
    df_compare["p_win"] = p_best
    threshold = np.quantile(p_best, 1 - args.top_decile)
    df_compare["take_top"] = (df_compare["p_win"] >= threshold).astype(int)
    df_compare["base_cum"] = df_compare["r_multiple"].cumsum()
    df_compare["top_cum"] = (df_compare["r_multiple"] * df_compare["take_top"]).cumsum()

    fig, ax = plt.subplots(figsize=(11, 5))
    ax.plot(df_compare["entry_time"], df_compare["base_cum"],
            label="Trade everything (engine baseline)", linewidth=1.5)
    ax.plot(df_compare["entry_time"], df_compare["top_cum"],
            label=f"Top {int(args.top_decile*100)}% by p_win (meta-model)",
            linewidth=1.8)
    ax.axhline(0, color="gray", linewidth=0.5)
    ax.set_xlabel("Trade entry time")
    ax.set_ylabel("Cumulative R")
    ax.set_title(f"Cumulative R: engine baseline vs meta-model top-{int(args.top_decile*100)}%  (test set)")
    ax.legend(); ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / "equity_compare.png", dpi=120)
    plt.close(fig)

    print(f"📁 Reports → {out_dir}/")
    print(f"   • classification_report.txt")
    print(f"   • feature_importances.txt")
    print(f"   • decile_breakdown.csv")
    print(f"   • summary.txt")
    print(f"📈 Plots:")
    print(f"   • calibration.png")
    print(f"   • decile_returns.png")
    print(f"   • equity_compare.png")
    return 0


if __name__ == "__main__":
    sys.exit(main())