#!/usr/bin/env python3
"""Generate a script that computes SHAP values correctly for a saved model and exports them as CSV.

The script it writes loads a joblib model (a bare estimator or a scikit-learn Pipeline) and a data CSV, then:
  - for a Pipeline, transforms the data with all steps but the last and explains the final estimator on the
    transformed columns (names from get_feature_names_out); with --group, sums the SHAP columns produced
    from one input column (one-hot levels, splines) back into that column, which keeps additivity exact
  - chooses the explainer from the model type at run time:
      LightGBM / XGBoost / sklearn trees and forests -> shap.TreeExplainer
          (--perturbation interventional with a background sample: SHAP relative to the background mean
           prediction; tree_path_dependent: no background, uses the tree's cover statistics; always used for
           HistGradientBoosting, whose interventional output in shap 0.54 misses the baseline)
      CatBoost  -> CatBoost's own get_feature_importance(type="ShapValues") (native categorical features)
      linear / logistic -> shap.LinearExplainer with the background
      anything else -> shap.PermutationExplainer on predict_proba / predict (model-agnostic, slower)
  - explains the raw model output (log-odds for boosted and linear classifiers, probability for sklearn
    forests and trees, the prediction for regressors) for tree and linear models, and the probability for the model-agnostic path; selects --class for multiclass
  - writes shap_values.csv (features + base_value + prediction) and shap_features.csv (the explained
    feature values) for shap_summary.py and explain_row.py, and asserts additivity
The generated script needs shap, pandas, numpy, joblib and the model's library.

Usage:
    python gen_shap_script.py --model model.joblib --data test.csv --target churned -o explain.py
    python gen_shap_script.py --model pipe.joblib --data test.csv --drop customer_id --group --rows 2000
    python gen_shap_script.py --model rf.joblib --data x.csv --perturbation tree_path_dependent --class 2

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

import argparse
import os
import sys

TEMPLATE = '''"""Compute SHAP values for {model!r} on {data!r}. Generated by gen_shap_script.py.

Run:  python {script}
Then: python shap_summary.py shap_values.csv --features shap_features.csv
      python explain_row.py shap_values.csv --features shap_features.csv --row 0{link_hint}
"""
import joblib
import numpy as np
import pandas as pd
import shap

MODEL, DATA, TARGET, DROP = {model!r}, {data!r}, {target!r}, {drop!r}
ROWS, BACKGROUND, PERTURBATION, CLASS, GROUP, SEED = {rows}, {background}, {perturbation!r}, {cls}, {group}, 0

model = joblib.load(MODEL)
df = pd.read_csv(DATA)
X_raw = df.drop(columns=[c for c in [TARGET] + DROP if c and c in df.columns])
if len(X_raw) > ROWS:
    X_raw = X_raw.sample(ROWS, random_state=SEED)
X_raw = X_raw.reset_index(drop=True)

steps = getattr(model, "steps", None)
if steps:                                           # explain the last step on the transformed data
    prep, est = model[:-1], model[-1]
    X = prep.transform(X_raw)
    X = X.toarray() if hasattr(X, "toarray") else np.asarray(X)
    names = list(prep.get_feature_names_out())
    X = pd.DataFrame(X, columns=names)
else:
    prep, est, X = None, model, X_raw.copy()
    names = list(X.columns)

kind = type(est).__module__.split(".")[0] + "." + type(est).__name__
is_clf = hasattr(est, "predict_proba")
background = X.sample(min(BACKGROUND, len(X)), random_state=SEED)


def pick(values, base):
    values, base = np.asarray(values), np.atleast_1d(np.asarray(base, dtype=float))
    if values.ndim == 3:                            # (rows, features, classes)
        k = CLASS if CLASS is not None else values.shape[2] - 1
        return values[:, :, k], float(base[k])
    return values, float(base[-1] if base.size > 1 and CLASS is None else base[CLASS or 0] if base.size > 1 else base[0])


if kind.startswith("catboost."):
    from catboost import Pool
    cat = [i for i, c in enumerate(X.columns) if not pd.api.types.is_numeric_dtype(X[c])]
    raw = est.get_feature_importance(Pool(X, cat_features=cat or None), type="ShapValues")
    if raw.ndim == 3:                               # (rows, classes, features + 1)
        raw = raw[:, CLASS if CLASS is not None else -1, :]
    sv, base = raw[:, :-1], float(raw[0, -1])
    output = est.predict(X, prediction_type="RawFormulaVal")
    output = output[:, CLASS if CLASS is not None else -1] if np.ndim(output) == 2 else output
    space = "raw (log-odds for classifiers)" if is_clf else "prediction"
elif kind.split(".")[0] in ("lightgbm", "xgboost") or "Tree" in kind or "Forest" in kind or "Boosting" in kind:
    pert = PERTURBATION
    if "HistGradientBoosting" in kind and pert == "interventional":
        # shap 0.54 interventional output for HistGradientBoosting misses the baseline (verified error ~2 log-odds)
        print("note: HistGradientBoosting -> tree_path_dependent (interventional is wrong for it in shap)")
        pert = "tree_path_dependent"
    if pert == "interventional":
        explainer = shap.TreeExplainer(est, data=background, feature_perturbation="interventional",
                                       model_output="raw")
    else:
        explainer = shap.TreeExplainer(est, feature_perturbation="tree_path_dependent")
    sv, base = pick(explainer.shap_values(X, check_additivity=True), explainer.expected_value)
    output = base + sv.sum(axis=1)
    forest = "Forest" in kind or "DecisionTree" in kind or "ExtraTrees" in kind
    space = ("probability" if forest else "raw (log-odds)") if is_clf else "prediction"
elif hasattr(est, "coef_"):
    explainer = shap.LinearExplainer(est, background)
    sv, base = pick(explainer.shap_values(X), explainer.expected_value)
    output = base + sv.sum(axis=1)
    space = "raw (log-odds)" if is_clf else "prediction"
else:
    if is_clf:
        k = CLASS if CLASS is not None else -1
        f = lambda a: est.predict_proba(pd.DataFrame(a, columns=names))[:, k]
    else:
        f = lambda a: est.predict(pd.DataFrame(a, columns=names))
    explainer = shap.PermutationExplainer(f, shap.maskers.Independent(background, max_samples=BACKGROUND),
                                          seed=SEED)
    e = explainer(X)
    sv, base = e.values, float(np.atleast_1d(e.base_values)[0])
    output = f(X.to_numpy())
    space = "probability" if is_clf else "prediction"

sv = np.asarray(sv, dtype=float)
err = np.abs(base + sv.sum(axis=1) - np.asarray(output, dtype=float)).max()
assert err < 1e-3 * max(1.0, float(np.abs(output).max())), f"additivity failed: max error {{err}}"

shap_df = pd.DataFrame(sv, columns=names)
feat_df = X.copy()
if GROUP and prep is not None:                      # sum SHAP of derived columns back to their input column
    originals = sorted(X_raw.columns, key=len, reverse=True)
    def source(name):
        short = name.split("__", 1)[-1]
        return next((o for o in originals if short == o or short.startswith(o + "_")), short)
    groups = {{}}
    for n in names:
        groups.setdefault(source(n), []).append(n)
    shap_df = pd.DataFrame({{g: shap_df[cols].sum(axis=1) for g, cols in groups.items()}})
    feat_df = X_raw[[g for g in groups if g in X_raw.columns]].copy()
    shap_df = shap_df[feat_df.columns]

shap_df["base_value"] = base
shap_df["prediction"] = output
shap_df.to_csv("shap_values.csv", index=False)
feat_df.to_csv("shap_features.csv", index=False)
print(f"{{kind}}: explained {{len(X)}} rows in {{space}} space; base value {{base:.5g}}; additivity error {{err:.2e}}")
print("wrote shap_values.csv and shap_features.csv")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a SHAP export 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="CSV with the rows to explain")
    ap.add_argument("--target", default="", help="target column to drop if present")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--rows", type=int, default=2000, help="rows to explain (sampled)")
    ap.add_argument("--background", type=int, default=200, help="background sample size")
    ap.add_argument("--perturbation", choices=["interventional", "tree_path_dependent"], default="interventional")
    ap.add_argument("--class", dest="cls", type=int, default=None, help="class index for multiclass models")
    ap.add_argument("--group", action="store_true", help="sum one-hot/derived SHAP columns per input column")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.rows < 1 or a.background < 10:
        print("error: --rows must be >= 1 and --background >= 10", file=sys.stderr)
        return 2
    text = TEMPLATE.format(model=a.model, data=a.data, target=a.target, drop=a.drop, rows=a.rows,
                           background=a.background, perturbation=a.perturbation, cls=a.cls, group=a.group,
                           script=os.path.basename(a.output or "explain.py"),
                           link_hint=" --link logit   (classifiers with log-odds output)")
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {a.perturbation} SHAP, {a.rows} rows, background {a.background}"
              + (", grouped by input column" if a.group else ""), file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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