#!/usr/bin/env python3
"""Generate a script that computes feature importance several ways on held-out data.

The script it writes loads a fitted model or Pipeline (joblib) and a HELD-OUT CSV (rows the model did not
train on), then computes, per input column:
  gain            impurity importance from the fitted model (LightGBM gain, XGBoost total_gain, CatBoost
                  PredictionValuesChange, sklearn feature_importances_), or coef_x_std (|coefficient| x
                  feature std) for linear models, mapped back from encoded columns to input columns by summing
  permutation     sklearn.inspection.permutation_importance on the whole pipeline with --repeats shuffles:
                  mean drop in the score and its std (so "indistinguishable from zero" can be judged)
  group perm.     joint permutation of clusters of correlated numeric features (|Spearman| >= --corr): the
                  importance of information that single-feature permutation hides when features share it
  mean_abs_shap   with --shap, mean |SHAP| from TreeExplainer (tree models), summed per input column
Held-out columns are reordered to the model's training order (CatBoost matches by position). It writes
importances.csv (wide: feature, gain or coef_x_std, permutation, permutation_std[, mean_abs_shap]) for
compare_importances.py, plus groups.csv with the grouped permutation results.
Needs scikit-learn, pandas, numpy, joblib (and shap with --shap).

Usage:
    python gen_importance_script.py --model model.joblib --data holdout.csv --target churned -o importance.py
    python gen_importance_script.py --model pipe.joblib --data test.csv --target price --scoring neg_root_mean_squared_error
    python gen_importance_script.py --model m.joblib --data test.csv --target y --shap --repeats 20 --corr 0.8

Exit codes: 0 written, 2 usage error.
"""

import argparse
import os
import sys

TEMPLATE = '''"""Feature importance for {model!r} on held-out {data!r}. Generated by gen_importance_script.py.

Run:  python {script}
Then: python compare_importances.py importances.csv
"""
import joblib
import numpy as np
import pandas as pd
from sklearn.inspection import permutation_importance
from sklearn.metrics import get_scorer

MODEL, DATA, TARGET, DROP = {model!r}, {data!r}, {target!r}, {drop!r}
SCORING, REPEATS, CORR, SHAP, SEED = {scoring!r}, {repeats}, {corr}, {shap}, 0

model = joblib.load(MODEL)
df = pd.read_csv(DATA)
df = df[df[TARGET].notna()]
X = df.drop(columns=[c for c in [TARGET] + DROP if c in df.columns])
y = df[TARGET]
trained = getattr(model, "feature_names_in_", None)
if trained is None:
    trained = getattr(model, "feature_names_", None)        # CatBoost matches columns by position
if trained is not None:
    missing = [c for c in trained if c not in X.columns]
    if missing:
        raise SystemExit(f"held-out data lacks training columns: {{missing}}")
    X = X[list(trained)]
if hasattr(model, "classes_") and y.dtype != object and set(np.unique(y)) - set(model.classes_):
    raise SystemExit("target values do not match model.classes_")
cols = list(X.columns)
scorer = get_scorer(SCORING)
baseline = scorer(model, X, y)
print(f"held-out {{SCORING}}: {{baseline:.5f}} on {{len(X)}} rows")

steps = getattr(model, "steps", None)
est = model[-1] if steps else model
names = list(model[:-1].get_feature_names_out()) if steps else cols


def to_input(name):
    """Map an encoded column name (e.g. cat__plan_basic) back to its input column."""
    short = name.split("__", 1)[-1]
    for c in sorted(cols, key=len, reverse=True):
        if short == c or short.startswith(c + "_"):
            return c
    return short


def per_input(values):
    out = dict.fromkeys(cols, 0.0)
    for n, v in zip(names, values):
        key = to_input(n)
        out[key] = out.get(key, 0.0) + float(v)
    return out


# 1) impurity / model-based importance
kind = type(est).__module__.split(".")[0]
label = "gain"
if kind == "lightgbm":
    raw = est.booster_.feature_importance(importance_type="gain")
elif kind == "xgboost":
    score = est.get_booster().get_score(importance_type="total_gain")
    fn = est.get_booster().feature_names or [f"f{{i}}" for i in range(len(names))]
    raw = [score.get(f, 0.0) for f in fn]
elif kind == "catboost":
    raw = est.get_feature_importance()
elif hasattr(est, "feature_importances_"):
    raw = est.feature_importances_
elif hasattr(est, "coef_"):
    Xt = model[:-1].transform(X) if steps else X.to_numpy()
    Xt = Xt.toarray() if hasattr(Xt, "toarray") else np.asarray(Xt, dtype=float)
    raw = np.abs(np.atleast_2d(est.coef_)).mean(axis=0) * Xt.std(axis=0)
    label = "coef_x_std"                                    # |coefficient| x feature std: not an impurity measure
else:
    raw = None
gain = per_input(raw) if raw is not None and len(raw) == len(names) else None

# 2) permutation importance on the whole pipeline (input columns)
perm = permutation_importance(model, X, y, scoring=SCORING, n_repeats=REPEATS, random_state=SEED, n_jobs=-1)

# 3) grouped permutation for clusters of correlated numeric features
num = [c for c in cols if pd.api.types.is_numeric_dtype(X[c])]
corr = X[num].corr(method="spearman").abs().fillna(0) if num else pd.DataFrame()
parent = {{c: c for c in num}}
def find(c):
    while parent[c] != c:
        parent[c] = parent[parent[c]]
        c = parent[c]
    return c
for i, a in enumerate(num):
    for b in num[i + 1:]:
        if corr.loc[a, b] >= CORR:
            parent[find(a)] = find(b)
clusters = {{}}
for c in num:
    clusters.setdefault(find(c), []).append(c)
rng = np.random.default_rng(SEED)
group_rows = []
for members in clusters.values():
    if len(members) < 2:
        continue
    drops = []
    for _ in range(REPEATS):
        Xp = X.copy()
        idx = rng.permutation(len(Xp))
        Xp[members] = Xp[members].to_numpy()[idx]           # permute the group jointly (keeps its internal structure)
        drops.append(baseline - scorer(model, Xp, y))
    single = sum(perm.importances_mean[cols.index(m)] for m in members)
    group_rows.append({{"group": "+".join(members), "size": len(members), "group_permutation": np.mean(drops),
                        "group_permutation_std": np.std(drops), "sum_single_permutation": single}})

# 4) optional mean |SHAP|
shap_imp = None
if SHAP:
    import shap
    Xt = model[:-1].transform(X) if steps else X
    Xt = Xt.toarray() if hasattr(Xt, "toarray") else Xt
    sv = shap.TreeExplainer(est).shap_values(Xt)
    sv = np.asarray(sv)
    if sv.ndim == 3:
        sv = sv[:, :, -1]
    shap_imp = per_input(np.abs(sv).mean(axis=0))

out = pd.DataFrame({{"feature": cols}})
if gain is not None:
    out[label] = [gain[c] for c in cols]
out["permutation"] = perm.importances_mean
out["permutation_std"] = perm.importances_std
if shap_imp is not None:
    out["mean_abs_shap"] = [shap_imp[c] for c in cols]
out = out.sort_values("permutation", ascending=False)
out.to_csv("importances.csv", index=False)
pd.DataFrame(group_rows).to_csv("groups.csv", index=False)
print(out.round(5).to_string(index=False))
for g in group_rows:
    print(f"group {{g['group']}}: joint permutation {{g['group_permutation']:.5f}} "
          f"vs sum of single {{g['sum_single_permutation']:.5f}}")
print("wrote importances.csv and groups.csv")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a feature importance script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", required=True, help="joblib file with the fitted model or pipeline")
    ap.add_argument("--data", required=True, help="held-out CSV (not used for training)")
    ap.add_argument("--target", required=True)
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--scoring", default="roc_auc", help="sklearn scorer name (default roc_auc)")
    ap.add_argument("--repeats", type=int, default=10)
    ap.add_argument("--corr", type=float, default=0.7, help="|Spearman| threshold for correlated groups")
    ap.add_argument("--shap", action="store_true", help="also compute mean |SHAP| (tree models)")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.repeats < 2 or not 0 < a.corr <= 1:
        print("error: --repeats must be >= 2 and 0 < --corr <= 1", file=sys.stderr)
        return 2
    text = TEMPLATE.format(model=a.model, data=a.data, target=a.target, drop=a.drop, scoring=a.scoring,
                           repeats=a.repeats, corr=a.corr, shap=a.shap,
                           script=os.path.basename(a.output or "importance.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {a.scoring}, {a.repeats} repeats, groups at |rho| >= {a.corr}"
              + (", with SHAP" if a.shap else ""), file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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