#!/usr/bin/env python3
"""
Compare prediction confidence between old (dev-main, 40 features) and
new (feat/chip-score, 63 features) model on the same live data.

Usage:
  python scripts/compare_versions.py
  python scripts/compare_versions.py --stocks 2330 0050 2317 2454
"""

import argparse
import os
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

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

LABEL_HORIZON = 5
LABEL_THRESHOLD = 0.02

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


def fetch_df(stock_no: str) -> pd.DataFrame:
    from services.predictor_service import _fetch_with_cache
    from indicators.technical import add_all_indicators
    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
    from indicators.technical import add_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


def predict_with_features(feat: pd.DataFrame, cols: list[str]) -> dict:
    avail = [c for c in cols if c in feat.columns]
    labels = feat["_label"].dropna()
    common = feat.index.intersection(labels.index)
    X = feat.loc[common, avail].fillna(0).values
    y = labels.loc[common].values

    # Drop last LABEL_HORIZON rows (no forward return yet)
    X_train, y_train = X[:-5], y[:-5]
    X_last = X[[-1]]

    if len(np.unique(y_train)) < 2:
        return {"signal": "N/A", "prob_up": float("nan"),
                "prob_dn": float("nan"), "prob_nt": float("nan"),
                "n_features": len(avail), "train_rows": len(X_train)}

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

    classes = list(clf.classes_)
    proba = clf.predict_proba(X_last)[0]
    prob_map = dict(zip(classes, proba))
    prob_up = prob_map.get(1, 0.0)
    prob_dn = prob_map.get(-1, 0.0)
    prob_nt = prob_map.get(0, 0.0)
    best = max(prob_up, prob_dn, prob_nt)
    signal = "UP" if prob_up == best else ("DOWN" if prob_dn == best else "HOLD")

    return {
        "signal": signal,
        "prob_up": round(prob_up * 100, 1),
        "prob_dn": round(prob_dn * 100, 1),
        "prob_nt": round(prob_nt * 100, 1),
        "n_features": len(avail),
        "train_rows": len(X_train),
    }


def predict_stock(stock_no: str):
    from models.predictor import _build_features

    df = fetch_df(stock_no)
    if df is None or df.empty:
        return None

    feat = _build_features(df)
    close = df.set_index("date")["close"] if "date" in df.columns else df["close"]
    fwd = close.shift(-LABEL_HORIZON)
    ret = (fwd - close) / close
    labels = pd.Series(
        np.where(ret > LABEL_THRESHOLD, 1, np.where(ret < -LABEL_THRESHOLD, -1, 0)),
        index=feat.index,
    ).astype(float)
    labels.iloc[-LABEL_HORIZON:] = np.nan
    feat["_label"] = labels

    return {
        f"old ({len(OLD_FEATURES)} feat)": predict_with_features(feat, OLD_FEATURES),
        f"new ({len(NEW_FEATURES)} feat)": predict_with_features(feat, NEW_FEATURES),
    }


def arrow(old_sig: str, new_sig: str) -> str:
    if old_sig == new_sig:
        return "="
    return f"{old_sig}→{new_sig}"


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

    hdr = f"{'Stock':>6}  {'Version':<18}  {'Signal':>5}  {'↑%':>6}  {'↓%':>6}  {'⬄%':>6}  {'Feat':>4}  {'Rows':>5}"
    sep = "-" * len(hdr)
    print(f"\n{hdr}\n{sep}")

    for stock_no in args.stocks:
        print(f"  fetching {stock_no}...", end="\r", flush=True)
        try:
            results = predict_stock(stock_no)
        except Exception as e:
            print(f"{stock_no:>6}  ERROR: {e}")
            continue

        if results is None:
            print(f"{stock_no:>6}  No data")
            continue

        versions = list(results.items())
        first = True
        for version, r in versions:
            prefix = f"{stock_no:>6}" if first else f"{'':>6}"
            first = False
            print(
                f"{prefix}  {version:<18}  {r['signal']:>5}  "
                f"{r['prob_up']:>5.1f}%  {r['prob_dn']:>5.1f}%  {r['prob_nt']:>5.1f}%  "
                f"{r['n_features']:>4}  {r.get('train_rows', 0):>5}"
            )

        # Summary delta line
        if len(versions) == 2:
            _, r_old = versions[0]
            _, r_new = versions[1]
            if r_old["signal"] != "N/A" and r_new["signal"] != "N/A":
                delta_up = r_new["prob_up"] - r_old["prob_up"]
                delta_dn = r_new["prob_dn"] - r_old["prob_dn"]
                sig_change = arrow(r_old["signal"], r_new["signal"])
                print(
                    f"{'':>6}  {'  delta':<18}  {sig_change:>5}  "
                    f"{delta_up:>+5.1f}%  {delta_dn:>+5.1f}%"
                )
        print()


if __name__ == "__main__":
    main()