#!/usr/bin/env python3
"""C31: SHAP re-pruning under C19 adaptive labels + C30 peer returns.

C5 pruned 11 features using SHAP under old fixed labels (29-feature set).
Since then C19 changed label quality (adaptive pt_sl) and C30 added peer return
features.  Feature importances may have shifted: some previously borderline
features may now be noise, and peer returns may have displaced some raw returns.

Run SHAP across all 12 EXTENDED_STOCKS using the current 31-feature set and
C19 adaptive labels.  Drop features with mean|SHAP| < 0.001 across stocks.
Validate pruned set through both 5-stock and 12-stock gates.
"""
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 json
import argparse
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier

from scripts.improvement_harness import (
    fetch_df, build_triple_barrier_labels, walk_forward, compute_metrics,
    CURRENT_FEATURES, DEFAULT_STOCKS, EXTENDED_STOCKS,
    PASS_DIR_ACC, PASS_UP_PREC, RF_PARAMS,
)
from models.predictor import _build_features

# Sector map for peer returns (mirrors backtest_c30.py)
SECTOR_MAP = {
    "semis":       ["2330", "2454", "2303"],
    "electronics": ["2317", "2382", "2308"],
    "financials":  ["2881", "2882", "2886"],
    "etfs":        ["0050", "0056"],
    "telecom":     ["2412"],
}
STOCK_SECTOR = {s: sec for sec, members in SECTOR_MAP.items() for s in members}
ALL_STOCKS = [s for members in SECTOR_MAP.values() for s in members]

SHAP_THRESHOLD = 0.001
SHAP_SAMPLE    = 150   # rows per stock for SHAP (larger than C5's 100 — more stable)


def _add_peer_features(feat: pd.DataFrame, df: pd.DataFrame,
                       stock_no: str, peer_close: dict) -> pd.DataFrame:
    sector = STOCK_SECTOR.get(stock_no, "")
    peers  = [p for p in SECTOR_MAP.get(sector, []) if p != stock_no and p in peer_close]
    if not peers or "date" not in df.columns:
        feat["peer_ret_1d"] = 0.0
        feat["peer_ret_5d"] = 0.0
        return feat
    date_col = df["date"].astype(str).reset_index(drop=True)
    r1_cols, r5_cols = [], []
    for p in peers:
        s = peer_close[p]
        r1 = s.pct_change(1).shift(1)
        r5 = s.pct_change(5).shift(1)
        r1_cols.append(date_col.map(r1.to_dict()).astype(float))
        r5_cols.append(date_col.map(r5.to_dict()).astype(float))
    feat = feat.reset_index(drop=True)
    feat["peer_ret_1d"] = pd.concat(r1_cols, axis=1).mean(axis=1).ffill().bfill().fillna(0.0).values
    feat["peer_ret_5d"] = pd.concat(r5_cols, axis=1).mean(axis=1).ffill().bfill().fillna(0.0).values
    return feat


def compute_shap_importances(stock_no: str, feat: pd.DataFrame,
                              labels: np.ndarray) -> dict[str, float]:
    try:
        import shap
    except ImportError:
        print("  shap not installed — run: pip install shap")
        return {}

    avail = [c for c in CURRENT_FEATURES if c in feat.columns]
    X = feat[avail].fillna(0).values
    valid = ~np.isnan(labels)
    X_v, y_v = X[valid][:-5], labels[valid][:-5].astype(int)
    if len(np.unique(y_v)) < 2 or len(X_v) < 50:
        return {}

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

    explainer = shap.TreeExplainer(clf)
    sample_idx = np.random.default_rng(42).choice(len(X_v), min(SHAP_SAMPLE, len(X_v)), replace=False)
    shap_vals = explainer.shap_values(X_v[sample_idx])

    sv = np.array(shap_vals)
    if sv.ndim == 3:
        mean_abs = np.mean(np.abs(sv), axis=(0, 2))
    elif sv.ndim == 2:
        mean_abs = np.mean(np.abs(sv), axis=0)
    else:
        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 _avg(results, key):
    vals = [m[key] for m in results if m and not np.isnan(m.get(key, float("nan")))]
    return round(float(np.mean(vals)), 1) if vals else float("nan")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--shap-only", action="store_true",
                        help="Print SHAP importances without running walk-forward")
    args = parser.parse_args()

    print("\n=== C31: SHAP re-pruning under C19 labels + C30 peer returns ===")
    print(f"  Base feature set: {len(CURRENT_FEATURES)} features")
    print(f"  SHAP threshold:   {SHAP_THRESHOLD}")
    print(f"  SHAP universe:    {len(EXTENDED_STOCKS)} stocks (12-stock)")

    # Load all stock data + peer closes
    print("\n  Loading stock data...")
    all_dfs: dict[str, pd.DataFrame] = {}
    for s in ALL_STOCKS:
        df = fetch_df(s)
        if df is not None and not df.empty:
            all_dfs[s] = df
    peer_close: dict[str, pd.Series] = {}
    for s, df in all_dfs.items():
        if "date" in df.columns:
            peer_close[s] = df.set_index("date")["close"].astype(float)
    print(f"  Loaded {len(all_dfs)} stocks")

    # Phase 1: SHAP importances on 12-stock universe
    print("\n  Computing SHAP importances (12 stocks × ~30s each)...")
    all_shap: dict[str, dict[str, float]] = {}
    for stock_no in EXTENDED_STOCKS:
        print(f"  SHAP {stock_no}...", end=" ", flush=True)
        df = all_dfs.get(stock_no)
        if df is None:
            print("no data"); continue
        feat = _build_features(df)
        feat = _add_peer_features(feat, df, stock_no, peer_close)
        close  = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        labels = build_triple_barrier_labels(close)
        imps   = compute_shap_importances(stock_no, feat, labels)
        all_shap[stock_no] = imps
        print(f"ok ({len(imps)} features)")

    # Aggregate
    agg_shap: dict[str, float] = {}
    for feat_name in CURRENT_FEATURES:
        vals = [all_shap[s][feat_name] for s in EXTENDED_STOCKS if feat_name in all_shap.get(s, {})]
        agg_shap[feat_name] = float(np.mean(vals)) if vals else 0.0

    sorted_feats = sorted(agg_shap.items(), key=lambda x: x[1], reverse=True)
    print("\n  Top 10 by mean|SHAP|:")
    for name, val in sorted_feats[:10]:
        print(f"    {name:<35} {val:.4f}")
    print("  Bottom 15 by mean|SHAP|:")
    for name, val in sorted_feats[-15:]:
        print(f"    {name:<35} {val:.5f}")

    low_feats = [f for f in CURRENT_FEATURES if agg_shap.get(f, 0) < SHAP_THRESHOLD]
    print(f"\n  Features with mean|SHAP| < {SHAP_THRESHOLD}: {low_feats}")

    if args.shap_only:
        out = ROOT / "docs/c31_shap_importances.json"
        out.write_text(json.dumps({"importances": agg_shap, "low_features": low_feats}, indent=2))
        print(f"\n  Saved: {out}")
        return 0

    if not low_feats:
        print("\n  No features below threshold — current set is already tight.")
        result = {
            "experiment": "C31",
            "shap_importances": agg_shap,
            "low_features": [],
            "pruned_feature_set": CURRENT_FEATURES,
            "n_removed": 0,
            "passed": False,
            "reason": "no features below threshold",
        }
        (ROOT / "docs/c31_result.json").write_text(json.dumps(result, indent=2))
        return 1

    PRUNED = [f for f in CURRENT_FEATURES if f not in low_feats]
    print(f"\n  Pruned set: {len(PRUNED)} features (removed {len(low_feats)})")

    # Phase 2: 5-stock walk-forward
    print(f"\n  Walk-forward [{len(DEFAULT_STOCKS)}-stock]...")
    base5, pruned5 = [], []
    for stock_no in DEFAULT_STOCKS:
        df = all_dfs.get(stock_no)
        if df is None:
            df = fetch_df(stock_no)
        if df is None: continue
        feat = _build_features(df)
        feat = _add_peer_features(feat, df, stock_no, peer_close)
        close  = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        labels = build_triple_barrier_labels(close)
        mb = walk_forward(feat, labels, CURRENT_FEATURES)
        mp = walk_forward(feat, labels, PRUNED)
        base5.append(mb); pruned5.append(mp)
        print(f"    {stock_no}: base dir={mb.get('dir_accuracy')}% ↑prec={mb.get('up_precision')}%"
              f"  →  pruned dir={mp.get('dir_accuracy')}% ↑prec={mp.get('up_precision')}%")

    base5_agg  = {"dir_accuracy": _avg(base5, "dir_accuracy"), "up_precision": _avg(base5, "up_precision")}
    pruned5_agg = {"dir_accuracy": _avg(pruned5, "dir_accuracy"), "up_precision": _avg(pruned5, "up_precision")}
    pass5 = pruned5_agg["dir_accuracy"] >= PASS_DIR_ACC and pruned5_agg["up_precision"] >= PASS_UP_PREC
    print(f"\n  5-stock baseline: dir={base5_agg['dir_accuracy']}%  ↑prec={base5_agg['up_precision']}%")
    print(f"  5-stock pruned:   dir={pruned5_agg['dir_accuracy']}%  ↑prec={pruned5_agg['up_precision']}%")
    print(f"  5-stock gate:     {'PASS ✓' if pass5 else 'FAIL ✗'}")

    # Phase 3: 12-stock walk-forward (always run to check degradation)
    print(f"\n  Walk-forward [{len(EXTENDED_STOCKS)}-stock]...")
    base12, pruned12 = [], []
    for stock_no in EXTENDED_STOCKS:
        df = all_dfs.get(stock_no)
        if df is None: continue
        feat = _build_features(df)
        feat = _add_peer_features(feat, df, stock_no, peer_close)
        close  = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        labels = build_triple_barrier_labels(close)
        mb = walk_forward(feat, labels, CURRENT_FEATURES)
        mp = walk_forward(feat, labels, PRUNED)
        base12.append(mb); pruned12.append(mp)
        print(f"    {stock_no}: base dir={mb.get('dir_accuracy')}% ↑prec={mb.get('up_precision')}%"
              f"  →  pruned dir={mp.get('dir_accuracy')}% ↑prec={mp.get('up_precision')}%")

    base12_agg  = {"dir_accuracy": _avg(base12, "dir_accuracy"), "up_precision": _avg(base12, "up_precision")}
    pruned12_agg = {"dir_accuracy": _avg(pruned12, "dir_accuracy"), "up_precision": _avg(pruned12, "up_precision")}
    pass12 = pruned12_agg["dir_accuracy"] >= PASS_DIR_ACC and pruned12_agg["up_precision"] >= PASS_UP_PREC
    print(f"\n  12-stock baseline: dir={base12_agg['dir_accuracy']}%  ↑prec={base12_agg['up_precision']}%")
    print(f"  12-stock pruned:   dir={pruned12_agg['dir_accuracy']}%  ↑prec={pruned12_agg['up_precision']}%")
    print(f"  12-stock gate:     {'PASS ✓' if pass12 else 'FAIL ✗'}")

    passed = pass5 and pass12
    print(f"\n  Final verdict: {'PASS ✓' if passed else 'FAIL ✗'}")
    print(f"  Removed features: {low_feats}")
    print(f"  New feature count: {len(PRUNED)}")

    result = {
        "experiment": "C31",
        "description": "SHAP re-pruning under C19 adaptive labels + C30 peer returns",
        "shap_threshold": SHAP_THRESHOLD,
        "shap_importances": dict(sorted_feats),
        "low_features": low_feats,
        "pruned_feature_set": PRUNED,
        "n_features_before": len(CURRENT_FEATURES),
        "n_features_after": len(PRUNED),
        "n_removed": len(low_feats),
        "aggregate": {
            "5stock_baseline": base5_agg,
            "5stock_pruned": pruned5_agg,
            "12stock_baseline": base12_agg,
            "12stock_pruned": pruned12_agg,
        },
        "passed": passed,
        "pass_gate": {"dir_accuracy": PASS_DIR_ACC, "up_precision": PASS_UP_PREC},
    }

    out = ROOT / "docs/c31_result.json"
    out.write_text(json.dumps(result, indent=2))
    print(f"\n  Saved: {out}")
    return 0 if passed else 1


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