#!/usr/bin/env python3
"""Generate a script that runs LIME correctly on tabular rows, repeatedly, and exports the runs.

The script it writes loads a fitted model or Pipeline (joblib), the TRAINING data (LIME takes feature
statistics and category frequencies from it) and the rows to explain, then:
  - encodes non-numeric columns as integer codes for LIME and decodes them inside predict_fn, so the
    model always receives the same DataFrame columns and dtypes it was trained on (Pipelines with
    OneHotEncoder work unchanged); missing numerics are filled with the training median for LIME only
  - builds lime.lime_tabular.LimeTabularExplainer with categorical_features / categorical_names,
    quartile discretisation (readable rules such as "support_calls > 3"; --no-discretize keeps continuous
    features as slopes, which usually fits smooth models better), the kernel width you choose
    (default 0.75 * sqrt(n_features), LIME's default) and --samples perturbations
  - explains each requested row --repeats times with different seeds, recording every feature weight,
    the surrogate's weighted R^2 (score), its prediction at the row (local_pred) and the model's
    prediction (model_pred), into lime_runs_row<N>.csv for lime_stability.py
  - saves the first run of each row as lime_row<N>.html
Classification explains the probability of --class (default: the last class); regression explains predict.
Needs lime, scikit-learn, pandas, numpy, joblib.

Usage:
    python gen_lime_script.py --model model.joblib --train train.csv --rows test.csv --target churned --row 0 5 -o lime_run.py
    python gen_lime_script.py --model pipe.joblib --train train.csv --rows test.csv --target price --row 3 --samples 10000
    python gen_lime_script.py --model m.joblib --train train.csv --rows test.csv --target y --kernel-width 2.0 --repeats 20

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

import argparse
import os
import sys

TEMPLATE = '''"""LIME explanations for {model!r}. Generated by gen_lime_script.py.

Run:  python {script}
Then: python lime_stability.py lime_runs_row<N>.csv
"""
import joblib
import numpy as np
import pandas as pd
from lime.lime_tabular import LimeTabularExplainer

MODEL, TRAIN, ROWS, TARGET, DROP = {model!r}, {train!r}, {rows!r}, {target!r}, {drop!r}
ROW_IDS, REPEATS, SAMPLES, TOP, KERNEL_WIDTH, CLASS, SEED = {row_ids!r}, {repeats}, {samples}, {top}, {kernel}, {cls}, 0
DISCRETIZE = {discretize}

model = joblib.load(MODEL)
train = pd.read_csv(TRAIN)
rows = pd.read_csv(ROWS)
drop = [c for c in [TARGET] + DROP if c]
Xtr = train.drop(columns=[c for c in drop if c in train.columns])
trained = getattr(model, "feature_names_in_", None)
if trained is None:
    trained = getattr(model, "feature_names_", None)
cols = list(trained) if trained is not None else list(Xtr.columns)
Xtr = Xtr[cols]
Xrows = rows[cols]
cat_cols = [c for c in cols if not pd.api.types.is_numeric_dtype(Xtr[c])]
cat_idx = [cols.index(c) for c in cat_cols]
levels = {{c: sorted(Xtr[c].fillna("__nan__").astype(str).unique()) for c in cat_cols}}
medians = Xtr[[c for c in cols if c not in cat_cols]].median()
dtypes = Xtr.dtypes


def encode(frame):
    out = frame.copy()
    for c in cat_cols:
        lv = {{v: i for i, v in enumerate(levels[c])}}
        out[c] = out[c].fillna("__nan__").astype(str).map(lambda v: lv.get(v, 0))
    num = [c for c in cols if c not in cat_cols]
    out[num] = out[num].fillna(medians)
    return out.to_numpy(dtype=float)


def decode(arr):
    frame = pd.DataFrame(arr, columns=cols)
    for c in cat_cols:
        codes = np.clip(np.rint(frame[c]).astype(int), 0, len(levels[c]) - 1)
        frame[c] = [None if levels[c][k] == "__nan__" else levels[c][k] for k in codes]
    for c in cols:
        if c not in cat_cols and pd.api.types.is_integer_dtype(dtypes[c]):
            frame[c] = np.rint(frame[c]).astype(dtypes[c])
    return frame


is_clf = hasattr(model, "predict_proba")
if is_clf:
    classes = list(model.classes_)
    label = CLASS if CLASS is not None else len(classes) - 1
    predict_fn = lambda a: model.predict_proba(decode(a))
else:
    label = None
    predict_fn = lambda a: model.predict(decode(a))

explainer = LimeTabularExplainer(
    encode(Xtr), feature_names=cols, categorical_features=cat_idx,
    categorical_names={{i: levels[c] for i, c in zip(cat_idx, cat_cols)}},
    class_names=[str(c) for c in classes] if is_clf else None,
    mode="classification" if is_clf else "regression",
    discretize_continuous=DISCRETIZE, discretizer="quartile", kernel_width=KERNEL_WIDTH, random_state=SEED)

for rid in ROW_IDS:
    x = encode(Xrows.iloc[[rid]])[0]
    model_pred = float(predict_fn(x[None, :])[0, label] if is_clf else predict_fn(x[None, :])[0])
    records = []
    for r in range(REPEATS):
        explainer.random_state = np.random.RandomState(SEED + r)     # a different perturbation sample per run
        kw = dict(num_features=TOP, num_samples=SAMPLES)
        if is_clf:
            kw["labels"] = (label,)
        exp = explainer.explain_instance(x, predict_fn, **kw)
        weights = exp.as_list(label=label) if is_clf else exp.as_list()
        local = exp.local_pred[0] if np.ndim(exp.local_pred) else exp.local_pred
        score = exp.score[label] if isinstance(exp.score, dict) else exp.score
        for feat, w in weights:
            records.append((r, feat, w, score, float(local), model_pred))
        if r == 0:
            exp.save_to_file(f"lime_row{{rid}}.html")
    out = pd.DataFrame(records, columns=["run", "feature", "weight", "score", "local_pred", "model_pred"])
    out.to_csv(f"lime_runs_row{{rid}}.csv", index=False)
    top = out.groupby("feature").weight.agg(["mean", "std"]).reindex(
        out.groupby("feature").weight.apply(lambda w: w.abs().mean()).sort_values(ascending=False).index)
    print(f"row {{rid}}: model {{model_pred:.4f}}, surrogate R^2 {{out.score.mean():.2f}}")
    print(top.head(TOP).round(4).to_string())
    print(f"wrote lime_runs_row{{rid}}.csv and lime_row{{rid}}.html")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a repeated-LIME script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", required=True)
    ap.add_argument("--train", required=True, help="training CSV (LIME statistics)")
    ap.add_argument("--rows", required=True, help="CSV containing the rows to explain")
    ap.add_argument("--target", default="")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--row", type=int, nargs="+", default=[0], help="0-based row indices in --rows")
    ap.add_argument("--repeats", type=int, default=10)
    ap.add_argument("--samples", type=int, default=5000, help="perturbations per explanation")
    ap.add_argument("--top", type=int, default=8, help="features per explanation")
    ap.add_argument("--kernel-width", type=float, default=None, help="default 0.75 * sqrt(n_features)")
    ap.add_argument("--class", dest="cls", type=int, default=None)
    ap.add_argument("--no-discretize", action="store_true",
                    help="continuous features stay continuous: higher fidelity for smooth models, less readable")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.repeats < 2 or a.samples < 500 or min(a.row) < 0:
        print("error: --repeats must be >= 2, --samples >= 500 and rows >= 0", file=sys.stderr)
        return 2
    text = TEMPLATE.format(model=a.model, train=a.train, rows=a.rows, target=a.target, drop=a.drop, row_ids=a.row,
                           repeats=a.repeats, samples=a.samples, top=a.top, kernel=a.kernel_width, cls=a.cls,
                           discretize=not a.no_discretize,
                           script=os.path.basename(a.output or "lime_run.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: rows {a.row}, {a.repeats} repeats x {a.samples} samples", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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