#!/usr/bin/env python3
"""
Walk-forward backtest: compare old (40-feat) vs new (current model) feature set.

Strategy:
  - Expanding window starting at MIN_TRAIN_ROWS rows
  - Step forward STEP_DAYS at a time
  - For each step: train both models on all data up to cutoff,
    predict on next STEP_DAYS rows, record accuracy
  - Label: 5-day forward return > +2% → 1 (UP), < -2% → -1 (DOWN), else 0 (HOLD)
  - Metrics: accuracy, directional precision (UP+DOWN only, ignore HOLD), long-only P&L

Usage:
  ./venv/bin/python3 scripts/backtest_compare.py
  ./venv/bin/python3 scripts/backtest_compare.py --stocks 2330 0050 2317 2454 2881
"""

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

try:
    from dotenv import load_dotenv
    load_dotenv(ROOT / ".env")
except ImportError:
    pass

import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# ── Feature sets ────────────────────────────────────────────────────────────

OLD_FEATURES = [
    "return_1d", "return_5d", "return_10d", "return_20d",
    "close_ma5_ratio", "close_ma20_ratio", "ma5_ma20_ratio", "ma20_ma60_ratio",
    "rsi",
    "macd_hist", "macd_signal_ratio", "macd_hist_norm", "macd_hist_delta_1d",
    "macd_hist_slope_3d", "macd_cross_up", "macd_cross_down", "macd_above_zero",
    "bb_pct_b", "k", "d", "volume_ratio",
    "atr_ratio", "high_low_ratio", "obv_trend",
    "close_ma60_ratio", "log_volume_ratio",
    "volume_zscore", "price_volume_div", "volatility_20d",
    "taiex_return_5d", "taiex_ma20_ratio", "usdtwd_return_5d",
    "foreign_net_vol_ratio", "trust_net_vol_ratio", "dealer_net_vol_ratio",
    "institutional_net_vol_ratio", "institutional_5d_net_vol_ratio",
    "institutional_20d_zscore", "foreign_trust_alignment", "institutional_streak",
]

# Always use the live FEATURE_COLUMNS from predictor so this script stays in sync
from models.predictor import FEATURE_COLUMNS as NEW_FEATURES

# ── Backtest parameters ──────────────────────────────────────────────────────

LABEL_HORIZON  = 5    # predict 5-day forward return
LABEL_THRESH   = 0.02 # ±2% to label UP/DOWN
MIN_TRAIN_ROWS = 240  # ~1 year of trading days before first prediction
STEP_DAYS      = 21   # step forward ~1 month at a time

RF_PARAMS = dict(n_estimators=200, max_depth=6, min_samples_leaf=10,
                 class_weight="balanced", random_state=42, n_jobs=-1)

# ── Data fetching ────────────────────────────────────────────────────────────

def fetch_df(stock_no: str) -> pd.DataFrame:
    from services.predictor_service import _fetch_with_cache
    from indicators.technical import add_all_indicators, add_cross_asset_tw
    from data.institutional_flow import add_institutional_flow
    from data.margin_flow import add_margin_flow
    from data.fetcher import is_us_ticker, fetch_cross_asset_tw

    df = _fetch_with_cache(stock_no, months=24)
    if df.empty:
        return df
    df = add_all_indicators(df)
    if not is_us_ticker(stock_no):
        df = add_institutional_flow(df, stock_no)
        df = add_margin_flow(df, stock_no)
        start, end = str(df["date"].min()), str(df["date"].max())
        result = fetch_cross_asset_tw(start, end)
        taiex, usdtwd = result[0], result[1]
        sox = result[2] if len(result) > 2 else None
        tnx = result[3] if len(result) > 3 else None
        df = add_cross_asset_tw(df, taiex, usdtwd, sox_close=sox, tnx_close=tnx)
    return df

# ── Backtest engine ──────────────────────────────────────────────────────────

def make_labels(close: pd.Series) -> pd.Series:
    fwd = close.shift(-LABEL_HORIZON)
    ret = (fwd - close) / close
    return pd.Series(
        np.where(ret > LABEL_THRESH, 1, np.where(ret < -LABEL_THRESH, -1, 0)),
        index=close.index,
    )


def walk_forward(feat: pd.DataFrame, feature_cols: list[str]) -> dict:
    """
    Expanding-window walk-forward.
    Returns dict of aggregated metrics.
    """
    avail = [c for c in feature_cols if c in feat.columns]
    labels = feat["_label"].values
    X_all  = feat[avail].fillna(0).values
    close  = feat["_close"].values
    n      = len(feat)

    y_true_all, y_pred_all = [], []
    # For P&L: +1 if predicted UP and actually went up, -1 if predicted DOWN and went down
    pnl_model, pnl_buy_hold = [], []

    cutoff = MIN_TRAIN_ROWS
    while cutoff + STEP_DAYS + LABEL_HORIZON <= n:
        X_train = X_all[:cutoff]
        y_train = labels[:cutoff]

        # Test window: next STEP_DAYS rows (skip last LABEL_HORIZON — no label yet)
        test_end = min(cutoff + STEP_DAYS, n - LABEL_HORIZON)
        X_test  = X_all[cutoff:test_end]
        y_test  = labels[cutoff:test_end]
        c_test  = close[cutoff:test_end]

        if len(np.unique(y_train)) < 2 or len(X_test) == 0:
            cutoff += STEP_DAYS
            continue

        clf = RandomForestClassifier(**RF_PARAMS)
        clf.fit(X_train, y_train)
        y_pred = clf.predict(X_test)

        y_true_all.extend(y_test.tolist())
        y_pred_all.extend(y_pred.tolist())

        # Simple directional P&L: 5-day fwd return when signal is UP or DOWN
        classes = list(clf.classes_)
        proba   = clf.predict_proba(X_test)
        prob_map = {cls: proba[:, i] for i, cls in enumerate(classes)}
        for i in range(len(X_test)):
            p_up = prob_map.get(1, np.zeros(len(X_test)))[i]
            p_dn = prob_map.get(-1, np.zeros(len(X_test)))[i]
            fwd_ret = (close[cutoff + i + LABEL_HORIZON] - c_test[i]) / c_test[i] if (cutoff + i + LABEL_HORIZON) < n else 0.0
            if y_pred[i] == 1:
                pnl_model.append(fwd_ret)
            elif y_pred[i] == -1:
                pnl_model.append(-fwd_ret)   # short position
            # Buy-and-hold baseline: always long
            pnl_buy_hold.append(fwd_ret)

        cutoff += STEP_DAYS

    if not y_true_all:
        return {}

    y_true = np.array(y_true_all)
    y_pred = np.array(y_pred_all)

    # Directional subset (exclude HOLD predictions for precision/recall)
    dir_mask = y_pred != 0
    y_true_dir = y_true[dir_mask]
    y_pred_dir = y_pred[dir_mask]

    acc = accuracy_score(y_true, y_pred)
    n_signals = dir_mask.sum()
    dir_acc   = accuracy_score(y_true_dir, y_pred_dir) if n_signals > 0 else float("nan")

    # UP-only precision: when model says UP, how often correct?
    up_mask = y_pred == 1
    up_prec = (y_true[up_mask] == 1).mean() if up_mask.sum() > 0 else float("nan")
    dn_mask = y_pred == -1
    dn_prec = (y_true[dn_mask] == -1).mean() if dn_mask.sum() > 0 else float("nan")

    total_pnl   = sum(pnl_model)
    bh_pnl      = sum(pnl_buy_hold)
    n_trades    = len(pnl_model)
    win_rate    = (np.array(pnl_model) > 0).mean() if pnl_model else float("nan")

    return {
        "n_predictions": len(y_true),
        "n_signals": int(n_signals),
        "accuracy": round(acc * 100, 1),
        "dir_accuracy": round(dir_acc * 100, 1) if not np.isnan(dir_acc) else float("nan"),
        "up_precision": round(up_prec * 100, 1) if not np.isnan(up_prec) else float("nan"),
        "dn_precision": round(dn_prec * 100, 1) if not np.isnan(dn_prec) else float("nan"),
        "pnl_pct": round(total_pnl * 100, 1),
        "bh_pnl_pct": round(bh_pnl * 100, 1),
        "win_rate": round(win_rate * 100, 1) if not np.isnan(win_rate) else float("nan"),
        "n_trades": n_trades,
    }


def fmt(val, suffix=""):
    if isinstance(val, float) and np.isnan(val):
        return "  N/A"
    return f"{val:>5.1f}{suffix}"


def run_stock(stock_no: str) -> dict | None:
    from models.predictor import _build_features

    df = fetch_df(stock_no)
    if df is None or df.empty or len(df) < MIN_TRAIN_ROWS + STEP_DAYS + LABEL_HORIZON:
        return None

    feat = _build_features(df)

    close_series = df.set_index("date")["close"] if "date" in df.columns else df["close"]
    labels = make_labels(close_series)
    feat["_label"] = labels.values
    feat["_close"] = close_series.values

    return {
        "old": walk_forward(feat, OLD_FEATURES),
        "new": walk_forward(feat, NEW_FEATURES),
    }


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--stocks", nargs="+", default=["2330", "0050", "2317", "2454", "2881"])
    args = parser.parse_args()

    print("\nWalk-forward backtest — expanding window, step=21d, label=5d ±2%")
    print(f"Min train: {MIN_TRAIN_ROWS} rows | Features: old={len(OLD_FEATURES)}, new={len(NEW_FEATURES)}\n")

    col_w = 54
    hdr = (f"{'Stock':<6}  {'Model':<5}  {'Acc%':>5}  {'Dir%':>5}  "
           f"{'↑Prec':>6}  {'↓Prec':>6}  {'WinR%':>6}  {'P&L%':>6}  {'B&H%':>6}  "
           f"{'Trades':>6}  {'Preds':>5}")
    print(hdr)
    print("-" * len(hdr))

    agg = {"old": [], "new": []}

    for stock_no in args.stocks:
        print(f"  computing {stock_no}...", end="\r", flush=True)
        try:
            results = run_stock(stock_no)
        except Exception as e:
            print(f"{stock_no:<6}  ERROR: {e}")
            continue
        if not results:
            print(f"{stock_no:<6}  insufficient data")
            continue

        for key in ("old", "new"):
            r = results[key]
            if not r:
                continue
            label = f"old{len(OLD_FEATURES)}" if key == "old" else f"new{len(NEW_FEATURES)}"
            print(
                f"{stock_no:<6}  {label:<5}  "
                f"{fmt(r['accuracy'])}  {fmt(r['dir_accuracy'])}  "
                f"{fmt(r['up_precision'])}%  {fmt(r['dn_precision'])}%  "
                f"{fmt(r['win_rate'])}%  {fmt(r['pnl_pct'])}%  "
                f"{fmt(r['bh_pnl_pct'])}%  {r['n_trades']:>6}  {r['n_predictions']:>5}"
            )
            agg[key].append(r)

        # Delta row
        ro, rn = results["old"], results["new"]
        if ro and rn:
            da = rn["accuracy"]     - ro["accuracy"]
            dd = rn["dir_accuracy"] - ro["dir_accuracy"]
            du = rn["up_precision"] - ro["up_precision"]
            dw = rn["win_rate"]     - ro["win_rate"]
            dp = rn["pnl_pct"]      - ro["pnl_pct"]
            print(
                f"{'':6}  {'Δ':<5}  "
                f"{da:>+5.1f}  {dd:>+5.1f}  "
                f"{du:>+5.1f}%  {'':>6}  "
                f"{dw:>+5.1f}%  {dp:>+5.1f}%"
            )
        print()

    # ── Aggregate summary ────────────────────────────────────────────────────
    if agg["old"] and agg["new"]:
        print("=" * len(hdr))
        print("AGGREGATE (mean across stocks):")
        for key, label in [("old", f"old{len(OLD_FEATURES)}"), ("new", f"new{len(NEW_FEATURES)}")]:
            rows = agg[key]
            def m(field):
                vals = [r[field] for r in rows if not np.isnan(r.get(field, float("nan")))]
                return round(sum(vals) / len(vals), 1) if vals else float("nan")
            print(
                f"  {label:<7}  acc={m('accuracy')}%  dir={m('dir_accuracy')}%  "
                f"↑prec={m('up_precision')}%  ↓prec={m('dn_precision')}%  "
                f"win={m('win_rate')}%  P&L={m('pnl_pct')}%  B&H={m('bh_pnl_pct')}%"
            )
        print()


if __name__ == "__main__":
    main()