"""
SHAP Feature Importance Analysis
=================================
Trains a LightGBM model on the sweep database and uses SHAP to rank which of
the 25 weight params actually drive strategy performance vs. which are noise.

Usage:
    python3 tools/shap_feature_importance.py                          # all assets/TFs
    python3 tools/shap_feature_importance.py --asset COINBASE_BTCUSD  # BTC only
    python3 tools/shap_feature_importance.py --asset COINBASE_BTCUSD --timeframe 6H
    python3 tools/shap_feature_importance.py --per-asset              # one table per asset

Output:
    Printed to stdout + written to results/shap_importance.md
"""

import argparse
import glob
import json
import os
import sqlite3
import sys
import warnings
from datetime import datetime

import lightgbm as lgb
import numpy as np
import pandas as pd
import shap

warnings.filterwarnings("ignore")

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import WEIGHT_COLS

DB_FILE = "results/sweep_database.db"
OUTPUT_FILE = "results/shap_importance.md"
# Pattern matches only the 20 active asset-TF files (e.g. *_COINBASE_BTCUSD_4H.json).
# Legacy template files are in strategies/params/templates/ and excluded by this pattern.
PARAMS_GLOB = "strategies/params/params_strategy_activation_scores_*_*.json"


def _load_param_lock_status():
    """Return dicts describing lock status of each WEIGHT_COL across all params files.

    Returns:
        disabled: set of param names locked to [0] in ALL params files that contain them.
        missing_some: set of param names absent from at least one params file.
    """
    files = glob.glob(PARAMS_GLOB)
    if not files:
        return set(), set()

    # For each weight col, collect values lists from every file that contains it.
    col_file_values: dict[str, list[list]] = {c: [] for c in WEIGHT_COLS}
    total_files = len(files)

    for path in files:
        try:
            with open(path) as f:
                cfg = json.load(f)
        except Exception:
            continue
        for col in WEIGHT_COLS:
            if col in cfg:
                entry = cfg[col]
                if "values" in entry:
                    col_file_values[col].append(entry["values"])
                else:
                    # range-based — not locked to 0
                    col_file_values[col].append(None)

    disabled = set()
    missing_some = set()

    for col in WEIGHT_COLS:
        present_count = len(col_file_values[col])
        if present_count < total_files:
            missing_some.add(col)
        if present_count == 0:
            continue
        # Disabled = every file that has the param locks it to [0]
        if all(v is not None and v == [0] for v in col_file_values[col]):
            disabled.add(col)

    return disabled, missing_some


def load_data(asset=None, timeframe=None):
    conn = sqlite3.connect(DB_FILE)
    q = "SELECT * FROM sweep_results WHERE calmar_ratio > 0"
    params = []
    if asset:
        q += " AND asset = ?"
        params.append(asset)
    if timeframe:
        q += " AND timeframe = ?"
        params.append(timeframe)
    df = pd.read_sql(q, conn, params=params)
    conn.close()
    return df


def run_shap(df, label="ALL", disabled=None, missing_some=None):
    disabled = disabled or set()
    missing_some = missing_some or set()

    cols = [c for c in WEIGHT_COLS if c in df.columns]
    X = df[cols].dropna()
    y = df.loc[X.index, "sortino_ratio"]

    # Drop rows where target is NaN or infinite
    mask = np.isfinite(y)
    X, y = X[mask], y[mask]

    if len(X) < 50:
        return None, None, f"  {label}: only {len(X)} rows — skipping (need ≥ 50)"

    model = lgb.LGBMRegressor(
        n_estimators=300,
        learning_rate=0.05,
        num_leaves=31,
        min_child_samples=20,
        verbose=-1,
    )
    model.fit(X, y)

    explainer = shap.TreeExplainer(model)
    shap_values = explainer.shap_values(X)

    mean_abs = pd.Series(np.abs(shap_values).mean(axis=0), index=cols)
    mean_signed = pd.Series(shap_values.mean(axis=0), index=cols)

    importance = pd.DataFrame({
        "mean_abs_shap": mean_abs,
        "mean_signed_shap": mean_signed,
    }).sort_values("mean_abs_shap", ascending=False)

    total = importance["mean_abs_shap"].sum()
    importance["pct_total"] = 100.0 * importance["mean_abs_shap"] / total
    importance["cumulative_pct"] = importance["pct_total"].cumsum()

    def verdict(row):
        param = row.name
        is_disabled = param in disabled
        is_missing = param in missing_some
        is_weak = row["pct_total"] < 2.0 and row["cumulative_pct"] > 80.0
        is_moderate = 1.0 <= row["pct_total"] < 2.0 and row["cumulative_pct"] > 80.0

        suffix = ""
        if is_missing:
            suffix = " ⚠️ MISSING FROM SOME PARAMS"

        if is_disabled:
            if is_weak:
                return f"REPLACE{suffix}"
            return f"DISABLED{suffix}"
        if is_weak:
            return f"WEAK ← consider locking to 0{suffix}"
        if is_moderate:
            return f"MODERATE{suffix}"
        return f"STRONG{suffix}"

    importance["verdict"] = importance.apply(verdict, axis=1)

    return importance, len(X), None


def format_table(importance, label, n_rows):
    lines = []
    lines.append(f"### {label}  ({n_rows:,} rows)")
    lines.append("")
    lines.append(f"| {'Rank':>4} | {'Param':<32} | {'Abs SHAP':>10} | {'Signed mean':>11} | {'% total':>7} | {'Cumul%':>7} | Verdict |")
    lines.append(f"|{'-'*6}|{'-'*34}|{'-'*12}|{'-'*13}|{'-'*9}|{'-'*9}|{'-'*30}|")
    for rank, (param, row) in enumerate(importance.iterrows(), 1):
        sign_str = f"{row['mean_signed_shap']:+.3f}"
        lines.append(
            f"| {rank:>4} | {param:<32} | {row['mean_abs_shap']:>10.3f} | "
            f"{sign_str:>11} | {row['pct_total']:>6.1f}% | "
            f"{row['cumulative_pct']:>6.1f}% | {row['verdict']} |"
        )
    lines.append("")
    weak = importance[importance["verdict"].str.startswith("WEAK")]
    if not weak.empty:
        lines.append(f"**WEAK params ({len(weak)}):** " +
                     ", ".join(f"`{p}`" for p in weak.index))
        lines.append("")
    return lines


def main():
    parser = argparse.ArgumentParser(
        description="SHAP feature importance from sweep DB",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--asset", help="Filter to specific asset (e.g. COINBASE_BTCUSD)")
    parser.add_argument("--timeframe", help="Filter to specific timeframe (e.g. 6H)")
    parser.add_argument("--per-asset", action="store_true",
                        help="Print one table per asset (all TFs combined)")
    parser.add_argument("--per-tf", action="store_true",
                        help="Print one table per asset × timeframe combination")
    args = parser.parse_args()

    if not os.path.exists(DB_FILE):
        print(f"DB not found: {DB_FILE}")
        sys.exit(1)

    all_lines = [
        "# SHAP Feature Importance",
        f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "DB: " + DB_FILE,
        "",
        "*Mean |SHAP| = average absolute contribution to Sortino prediction.*",
        "*Signed mean shows direction — positive = param pushes Sortino higher.*",
        "*WEAK = < 2% of total importance — candidate for locking to 0.*",
        "",
        "---",
        "",
    ]

    disabled, missing_some = _load_param_lock_status()

    if args.per_tf:
        df_all = load_data(asset=args.asset, timeframe=args.timeframe)
        combos = sorted(df_all[["asset", "timeframe"]].drop_duplicates().itertuples(index=False))
        for row in combos:
            asset, tf = row.asset, row.timeframe
            sub = df_all[(df_all["asset"] == asset) & (df_all["timeframe"] == tf)]
            label = f"{asset} / {tf}"
            imp, n, err = run_shap(sub, label, disabled=disabled, missing_some=missing_some)
            if err:
                print(err)
                all_lines.append(err)
                continue
            table = format_table(imp, label, n)
            for line in table:
                print(line)
            all_lines.extend(table)
    elif args.per_asset:
        df_all = load_data(asset=args.asset, timeframe=args.timeframe)
        assets = sorted(df_all["asset"].unique())
        for asset in assets:
            sub = df_all[df_all["asset"] == asset]
            imp, n, err = run_shap(sub, asset, disabled=disabled, missing_some=missing_some)
            if err:
                print(err)
                all_lines.append(err)
                continue
            table = format_table(imp, asset, n)
            for line in table:
                print(line)
            all_lines.extend(table)
    else:
        label = " + ".join(filter(None, [args.asset, args.timeframe])) or "ALL assets & timeframes"
        df = load_data(asset=args.asset, timeframe=args.timeframe)
        if df.empty:
            print("No data found.")
            sys.exit(1)
        print(f"Loaded {len(df):,} rows from DB")
        imp, n, err = run_shap(df, label, disabled=disabled, missing_some=missing_some)
        if err:
            print(err)
            sys.exit(1)
        table = format_table(imp, label, n)
        for line in table:
            print(line)
        all_lines.extend(table)

    all_lines.append("---")
    all_lines.append(f"*Generated by tools/shap_feature_importance.py*")

    os.makedirs("results", exist_ok=True)
    with open(OUTPUT_FILE, "w") as f:
        f.write("\n".join(all_lines) + "\n")
    print(f"\n[Written to {OUTPUT_FILE}]")


if __name__ == "__main__":
    main()