#!/usr/bin/env python3
"""C5 validation: SHAP-based feature pruning."""
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
import json
from sklearn.ensemble import RandomForestClassifier

sys.path.insert(0, str(ROOT / "scripts"))
from improvement_harness import (
    BASELINE_FEATURES, fetch_df, build_triple_barrier_labels,
    run_comparison, RF_PARAMS,
)
from models.predictor import _build_features

STOCKS = ["2330", "0050", "2317", "2454", "2881"]
SHAP_THRESHOLD = 0.001   # features with mean|SHAP| below this on ALL stocks → pruned


def compute_shap_importances(stock_no: str) -> dict[str, float]:
    """Return {feature: mean_abs_shap} for BASELINE_FEATURES on this stock."""
    try:
        import shap
    except ImportError:
        return {}

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

    feat = _build_features(df)
    avail = [c for c in BASELINE_FEATURES if c in feat.columns]
    X = feat[avail].fillna(0).values
    close = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
    labels = build_triple_barrier_labels(close)
    valid = ~np.isnan(labels)
    X_v, y_v = X[valid][:-5], labels[valid][:-5].astype(int)  # leave last 5 for test

    if len(np.unique(y_v)) < 2 or len(X_v) < 50:
        return {}

    clf = RandomForestClassifier(**RF_PARAMS)
    clf.fit(X_v, y_v)

    # Use shap.TreeExplainer — fast for RF
    explainer = shap.TreeExplainer(clf)
    # Compute on a sample of 100 rows to keep it fast
    sample_idx = np.random.default_rng(42).choice(len(X_v), min(100, len(X_v)), replace=False)
    shap_vals = explainer.shap_values(X_v[sample_idx])  # list of arrays, one per class

    # Aggregate: mean |SHAP| across all classes and samples
    # shap >= 0.46 returns ndarray (n_samples, n_features, n_classes) for multi-class RF
    # older versions return list of (n_samples, n_features)
    sv = np.array(shap_vals)
    if sv.ndim == 3:
        # (n_samples, n_features, n_classes) — average over samples and classes
        mean_abs = np.mean(np.abs(sv), axis=(0, 2))
    elif sv.ndim == 2:
        mean_abs = np.mean(np.abs(sv), axis=0)
    else:
        # list of arrays (old API): list[(n_samples, n_features)]
        mean_abs = np.mean([np.mean(np.abs(np.array(s)), axis=0) for s in shap_vals], axis=0)

    return {feat_name: float(mean_abs[i]) for i, feat_name in enumerate(avail)}


def main():
    # Phase 1: compute SHAP importances per stock
    print("Computing SHAP importances (this takes ~2 min)...")
    all_shap: dict[str, dict[str, float]] = {}
    for stock_no in STOCKS:
        print(f"  SHAP for {stock_no}...", end="\r", flush=True)
        all_shap[stock_no] = compute_shap_importances(stock_no)
        print(f"  SHAP for {stock_no} done ({len(all_shap[stock_no])} features)")

    # Aggregate: mean |SHAP| across stocks for each feature
    agg_shap: dict[str, float] = {}
    for feat_name in BASELINE_FEATURES:
        vals = [all_shap[s][feat_name] for s in STOCKS if feat_name in all_shap.get(s, {})]
        agg_shap[feat_name] = float(np.mean(vals)) if vals else 0.0

    # Sort and display
    sorted_feats = sorted(agg_shap.items(), key=lambda x: x[1], reverse=True)
    print("\nTop 10 features by mean |SHAP|:")
    for name, val in sorted_feats[:10]:
        print(f"  {name:<35} {val:.4f}")
    print("\nBottom 10 features by mean |SHAP|:")
    for name, val in sorted_feats[-10:]:
        print(f"  {name:<35} {val:.4f}")

    # Features with SHAP < threshold on all stocks → candidates to prune
    zero_feats = [f for f in BASELINE_FEATURES if agg_shap.get(f, 0) < SHAP_THRESHOLD]
    print(f"\nFeatures with mean|SHAP| < {SHAP_THRESHOLD}: {zero_feats}")

    if not zero_feats:
        print("No features to prune — all above threshold. Exiting.")
        result = {
            "shap_importances": agg_shap, "zero_features": [],
            "shap_pruned_set": BASELINE_FEATURES,
            "passed": False, "reason": "no features below threshold",
        }
        Path("docs").mkdir(exist_ok=True)
        with open("docs/c5_shap_result.json", "w") as f:
            json.dump(result, f, indent=2)
        return

    SHAP_PRUNED = [f for f in BASELINE_FEATURES if f not in zero_feats]
    print(f"\nSHAP_PRUNED set: {len(SHAP_PRUNED)} features (removed {len(zero_feats)})")

    # Phase 2: walk-forward comparison
    print("\nRunning walk-forward comparison...")
    cmp = run_comparison(
        feature_sets={"baseline": BASELINE_FEATURES, "shap_pruned": SHAP_PRUNED},
        label_mode="triple_barrier",
        output_path=None,  # we'll write our own JSON
        pass_criterion={"dir_accuracy": 41.2},
        verbose=True,
    )

    b = cmp["aggregate"].get("baseline", {})
    p = cmp["aggregate"].get("shap_pruned", {})
    no_regress  = (p.get("dir_accuracy", 0) >= b.get("dir_accuracy", 0) - 0.5 and
                   p.get("up_precision", 0) >= b.get("up_precision", 0) - 0.5)
    improvement = (p.get("dir_accuracy", 0) - b.get("dir_accuracy", 0) >= 1.0 or
                   p.get("up_precision", 0) - b.get("up_precision", 0) >= 1.0)
    passed = bool(no_regress and improvement)

    result = {
        "shap_importances": agg_shap,
        "zero_features": zero_feats,
        "shap_pruned_set": SHAP_PRUNED,
        "aggregate": cmp["aggregate"],
        "results": cmp["results"],
        "passed": passed,
        "pass_criterion": "no_regress AND >=1.0pp on either metric",
    }
    Path("docs").mkdir(exist_ok=True)
    with open("docs/c5_shap_result.json", "w") as f:
        json.dump(result, f, indent=2)

    print(f"\n  Pass: {'YES' if passed else 'NO'}")


if __name__ == "__main__":
    main()