#!/usr/bin/env python3
"""Generate a script that draws PDP + ICE (and ALE) plots for a saved model and exports the curves.

The script it writes loads a fitted model or Pipeline (joblib) and a data CSV and, for each --features column:
  - computes ICE curves with sklearn.inspection.partial_dependence(kind="individual") on a --sample of rows,
    a grid between the 5th and 95th percentiles (--grid points; categorical features use their levels) and
    the right response (predict_proba of the positive / --class for classifiers, predict for regressors)
  - writes ice_<feature>.csv (row_id, <feature>, prediction) and rows_<feature>.csv (the sampled rows with a
    row_id) for ice_analysis.py
  - draws pdp_<feature>.png: centred ICE (thin lines), the PDP (thick) and a rug of observed values
  - with --ale, computes first-order accumulated local effects (ALE) by quantile bins: the mean change in
    prediction across each bin for the rows inside it, accumulated at the bin edges and centred (exact for a
    linear model: the slope equals the coefficient), drawn in a panel below the ICE plot. ALE stays on realistic rows
    when features are correlated, where PDP averages over impossible combinations
  - with --pair a b, draws a two-way PDP heatmap (interaction of two features)
Needs scikit-learn >= 1.3, pandas, numpy, joblib, matplotlib.

Usage:
    python gen_pdp_script.py --model model.joblib --data test.csv --target churned --features tenure_months support_calls -o pdp.py
    python gen_pdp_script.py --model pipe.joblib --data test.csv --features sqft --ale --sample 500
    python gen_pdp_script.py --model m.joblib --data x.csv --features age --pair age income --class 1

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

import argparse
import os
import sys

TEMPLATE = '''"""PDP / ICE{ale_title} for {model!r} on {data!r}. Generated by gen_pdp_script.py.

Run:  python {script}
Then: python ice_analysis.py ice_<feature>.csv --features rows_<feature>.csv
"""
import joblib
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.inspection import partial_dependence

MODEL, DATA, TARGET, DROP = {model!r}, {data!r}, {target!r}, {drop!r}
FEATURES, PAIR, ALE = {features!r}, {pair!r}, {ale}
SAMPLE, GRID, CLASS, SEED = {sample}, {grid}, {cls}, 0

model = joblib.load(MODEL)
df = pd.read_csv(DATA)
X = df.drop(columns=[c for c in [TARGET] + DROP if c and c in df.columns])
trained = getattr(model, "feature_names_in_", None)
if trained is not None:
    X = X[list(trained)]
missing = [f for f in FEATURES + PAIR if f not in X.columns]
if missing:
    raise SystemExit(f"features not in the data: {{missing}}")
ints = [c for c in X.columns if pd.api.types.is_integer_dtype(X[c])]
X[ints] = X[ints].astype(float)                        # sklearn >= 1.6 refuses integer columns for PDP
Xs = X.sample(min(SAMPLE, len(X)), random_state=SEED).reset_index(drop=True)
is_clf = hasattr(model, "predict_proba")
CATEGORICAL = [c for c in X.columns if not pd.api.types.is_numeric_dtype(X[c])]


def predict(frame):
    if is_clf:
        p = model.predict_proba(frame)
        return p[:, CLASS if CLASS is not None else p.shape[1] - 1]
    return model.predict(frame)


def ice(feature):
    kw = dict(kind="individual", grid_resolution=GRID, percentiles=(0.05, 0.95),
              categorical_features=[feature] if feature in CATEGORICAL else None,
              response_method="predict_proba" if is_clf else "auto")
    if is_clf and CLASS is not None and len(getattr(model, "classes_", [])) > 2:
        r = partial_dependence(model, Xs, [feature], **kw)
        curves = r["individual"][CLASS]
    else:
        r = partial_dependence(model, Xs, [feature], **kw)
        curves = r["individual"][0]
    return r["grid_values"][0], curves


def ale(feature, bins=20):
    x = X[feature].to_numpy(dtype=float)
    edges = np.unique(np.quantile(x[~np.isnan(x)], np.linspace(0, 1, bins + 1)))
    effects, counts = [], []
    for k, (lo, hi) in enumerate(zip(edges[:-1], edges[1:])):
        inside = X[((X[feature] > lo) | (k == 0)) & (X[feature] >= lo) & (X[feature] <= hi)]   # first bin includes its lower edge
        up, down = inside.copy(), inside.copy()
        up[feature], down[feature] = hi, lo
        effects.append(float(np.mean(predict(up) - predict(down))) if len(inside) else 0.0)
        counts.append(len(inside))
    acc = np.concatenate([[0.0], np.cumsum(effects)])  # ALE is defined at the bin edges
    mids = (acc[:-1] + acc[1:]) / 2
    acc -= np.average(mids, weights=np.maximum(counts, 1e-12))   # centre: mean effect over the data is zero
    return edges, acc


for f in FEATURES:
    grid, curves = ice(f)
    rows = [(i, g, curves[i, k]) for i in range(curves.shape[0]) for k, g in enumerate(grid)]
    pd.DataFrame(rows, columns=["row_id", f, "prediction"]).to_csv(f"ice_{{f}}.csv", index=False)
    Xs.assign(row_id=range(len(Xs))).to_csv(f"rows_{{f}}.csv", index=False)
    categorical = f in CATEGORICAL
    with_ale = ALE and not categorical
    if with_ale:
        fig, (ax, ax2) = plt.subplots(2, 1, figsize=(6, 6.5), sharex=True, gridspec_kw={{"height_ratios": [3, 2]}})
    else:
        fig, ax = plt.subplots(figsize=(6, 4))
    xs = np.arange(len(grid)) if categorical else np.asarray(grid, dtype=float)
    centred = curves - curves[:, [0]]
    for c in centred[: min(200, len(centred))]:
        ax.plot(xs, c, color="#4C78A8", alpha=0.15, lw=0.8)
    ax.plot(xs, centred.mean(axis=0), color="#E45756", lw=2.5, label="PDP (centred)")
    if categorical:
        ax.set_xticks(xs)
        ax.set_xticklabels([str(g) for g in grid], rotation=30)
    else:
        ax.plot(X[f].dropna().sample(min(500, X[f].notna().sum()), random_state=SEED),
                np.full(min(500, X[f].notna().sum()), ax.get_ylim()[0]), "|", color="k", alpha=0.3)
    ax.set_xlabel("" if with_ale else f)
    ax.set_ylabel("change in " + ("predicted probability" if is_clf else "prediction"))
    ax.set_title(f"Centred ICE and PDP: {{f}}")
    ax.legend(loc="best")
    if with_ale:
        cx, acc = ale(f)
        ax2.plot(cx, acc, color="#54A24B", lw=2, marker=".", label="ALE (centred)")
        ax2.axhline(0, color="grey", lw=0.8)
        ax2.set_ylabel("ALE: " + ("probability" if is_clf else "prediction"))
        ax2.set_xlabel(f)
        ax2.legend(loc="best")
        pd.DataFrame({{f: cx, "ale": acc}}).to_csv(f"ale_{{f}}.csv", index=False)
    fig.tight_layout()
    fig.savefig(f"pdp_{{f}}.png", dpi=130)
    plt.close(fig)
    print(f"{{f}}: wrote ice_{{f}}.csv, rows_{{f}}.csv, pdp_{{f}}.png" + (f", ale_{{f}}.csv" if ALE and not categorical else ""))

if PAIR:
    a, b = PAIR
    r = partial_dependence(model, Xs, [a, b], kind="average",           # a list of two = joint (two-way) PDP grid_resolution=min(GRID, 20),
                           percentiles=(0.05, 0.95), response_method="predict_proba" if is_clf else "auto",
                           categorical_features=[c for c in (a, b) if c in CATEGORICAL] or None)
    z = r["average"][CLASS if (is_clf and CLASS is not None and r["average"].shape[0] > 1) else 0]
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(z.T, origin="lower", aspect="auto", cmap="viridis")
    ax.set_xticks(range(len(r["grid_values"][0]))[::3])
    ax.set_xticklabels([f"{{v:.3g}}" if isinstance(v, (int, float, np.floating)) else str(v)
                        for v in r["grid_values"][0]][::3])
    ax.set_yticks(range(len(r["grid_values"][1]))[::3])
    ax.set_yticklabels([f"{{v:.3g}}" if isinstance(v, (int, float, np.floating)) else str(v)
                        for v in r["grid_values"][1]][::3])
    ax.set_xlabel(a)
    ax.set_ylabel(b)
    fig.colorbar(im, label="partial dependence")
    ax.set_title(f"Two-way PDP: {{a}} x {{b}}")
    fig.tight_layout()
    fig.savefig(f"pdp2_{{a}}_{{b}}.png", dpi=130)
    print(f"wrote pdp2_{{a}}_{{b}}.png")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a PDP/ICE/ALE script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--target", default="")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--features", nargs="+", required=True)
    ap.add_argument("--pair", nargs=2, metavar=("A", "B"), help="two-way PDP for this pair")
    ap.add_argument("--ale", action="store_true", help="add accumulated local effects (numeric features)")
    ap.add_argument("--sample", type=int, default=300, help="rows for ICE curves")
    ap.add_argument("--grid", type=int, default=20, help="grid points per feature")
    ap.add_argument("--class", dest="cls", type=int, default=None, help="class index (default: last class)")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.sample < 20 or a.grid < 3:
        print("error: --sample must be >= 20 and --grid >= 3", file=sys.stderr)
        return 2
    text = TEMPLATE.format(model=a.model, data=a.data, target=a.target, drop=a.drop, features=a.features,
                           pair=a.pair or [], ale=a.ale, sample=a.sample, grid=a.grid, cls=a.cls,
                           ale_title=" / ALE" if a.ale else "", script=os.path.basename(a.output or "pdp.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {', '.join(a.features)}" + (" + ALE" if a.ale else "")
              + (f", pair {a.pair[0]} x {a.pair[1]}" if a.pair else ""), file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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