#!/usr/bin/env python3
"""Generate a DiCE script that finds diverse, actionable counterfactuals for any scikit-learn model.

The script it writes loads a fitted classifier or Pipeline (joblib), the training data (DiCE learns
feature ranges and categories from it) and the rows to explain, then for each row:
  - turns the constraints file into DiCE arguments: immutable features are left out of features_to_vary,
    direction constraints become a per-row permitted_range ([min, current] for "decrease",
    [current, max] for "increase"), min/max bounds are kept
  - calls dice_ml.Dice(...).generate_counterfactuals with --method (random | genetic | kdtree),
    --total counterfactuals, desired_class="opposite" (or --desired-class) and a fixed seed
  - rounds integer features, drops duplicates, re-predicts every counterfactual with the model (DiCE's
    output is re-checked, not trusted), and writes row<N>.json (the original) and cfs_row<N>.csv
    (counterfactuals + pred_prob) ready for cf_validate.py
Needs dice-ml, scikit-learn, pandas, joblib.

Usage:
    python gen_dice_script.py --model model.joblib --train train.csv --rows test.csv --target churned --row 0 -o dice_run.py
    python gen_dice_script.py --model m.joblib --train train.csv --rows test.csv --target churned \\
        --constraints constraints.json --method genetic --total 6 --row 3 7

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

import argparse
import os
import sys

TEMPLATE = '''"""DiCE counterfactuals for {model!r}. Generated by gen_dice_script.py.

Run:  python {script}
Then: run the cf_validate.py command printed for each row (it includes the right --side)
"""
import json

import dice_ml
import joblib
import numpy as np
import pandas as pd

MODEL, TRAIN, ROWS, TARGET, DROP = {model!r}, {train!r}, {rows!r}, {target!r}, {drop!r}
CONSTRAINTS, ROW_IDS, METHOD, TOTAL, DESIRED, SEED = {constraints!r}, {row_ids!r}, {method!r}, {total}, {desired!r}, 0

model = joblib.load(MODEL)
train = pd.read_csv(TRAIN)
rows = pd.read_csv(ROWS)
feats = [c for c in train.columns if c != TARGET and c not in DROP]
trained = getattr(model, "feature_names_in_", None)
if trained is not None:
    feats = list(trained)
train = train[feats + [TARGET]].dropna()
cons = json.load(open(CONSTRAINTS))["features"] if CONSTRAINTS else {{}}
continuous = [c for c in feats if pd.api.types.is_numeric_dtype(train[c]) and not cons.get(c, {{}}).get("binary")
              and train[c].nunique() > 2]
categorical = [c for c in feats if c not in continuous]
train[categorical] = train[categorical].astype(str)            # DiCE treats categoricals as strings


class Wrapped:
    """Give the model the dtypes it was trained on (DiCE passes categoricals as strings)."""
    def __init__(self, m, dtypes):
        self.m, self.dtypes = m, dtypes
    def _fix(self, X):
        X = X[feats].copy()
        for c in categorical:
            if pd.api.types.is_numeric_dtype(self.dtypes[c]):
                X[c] = pd.to_numeric(X[c]).astype(self.dtypes[c])
        return X
    def predict(self, X):
        return self.m.predict(self._fix(X))
    def predict_proba(self, X):
        return self.m.predict_proba(self._fix(X))


dtypes = pd.read_csv(TRAIN)[feats].dtypes
wrapped = Wrapped(model, dtypes)
data = dice_ml.Data(dataframe=train, continuous_features=continuous, outcome_name=TARGET)
dmodel = dice_ml.Model(model=wrapped, backend="sklearn", model_type="classifier")
explainer = dice_ml.Dice(data, dmodel, method=METHOD)

for rid in ROW_IDS:
    query = rows[feats].iloc[[rid]].copy()
    query[categorical] = query[categorical].astype(str)
    vary = [c for c in feats if cons.get(c, {{}}).get("mutable", True)]
    permitted = {{}}
    for c in continuous:
        cfg = cons.get(c, {{}})
        lo = cfg.get("min", float(train[c].min()))
        hi = cfg.get("max", float(train[c].max()))
        x = float(query[c].iloc[0])
        if cfg.get("direction") == "decrease":
            hi = min(hi, x)
        elif cfg.get("direction") == "increase":
            lo = max(lo, x)
        if c in vary and lo < hi:
            permitted[c] = [lo, hi]
        elif c in vary:
            vary.remove(c)                                         # no room to move in the allowed direction
    kw = {{"random_seed": SEED}} if METHOD == "random" else {{}}  # only the random method takes a seed argument
    np.random.seed(SEED)
    result = explainer.generate_counterfactuals(
        query, total_CFs=TOTAL, desired_class=DESIRED, features_to_vary=vary,
        permitted_range=permitted or None, **kw)
    cfs = result.cf_examples_list[0].final_cfs_df
    original = query.iloc[0].to_dict()
    with open(f"row{{rid}}.json", "w") as f:
        json.dump({{k: (v.item() if hasattr(v, "item") else v) for k, v in original.items()}}, f, indent=2)
    if cfs is None or cfs.empty:
        print(f"row {{rid}}: DiCE found no counterfactual within the constraints")
        continue
    cfs = cfs[feats].copy()
    for c in continuous:
        if cons.get(c, {{}}).get("integer") or pd.api.types.is_integer_dtype(dtypes[c]):
            cfs[c] = np.rint(cfs[c].astype(float)).astype(int)
    cfs = cfs.drop_duplicates()
    proba = wrapped.predict_proba(cfs)
    cfs["pred_prob"] = proba[:, -1]
    cfs.to_csv(f"cfs_row{{rid}}.csv", index=False)
    p0 = wrapped.predict_proba(query)[0, -1]
    side = "below" if p0 >= 0.5 else "above"                    # "opposite" flips the predicted class
    print(f"row {{rid}}: original probability {{p0:.3f}}; {{len(cfs)}} counterfactual(s) -> cfs_row{{rid}}.csv")
    print(f"validate: python cf_validate.py --original row{{rid}}.json --cfs cfs_row{{rid}}.csv --train {{TRAIN}}"
          + (f" --constraints {{CONSTRAINTS}}" if CONSTRAINTS else "") + f" --pred-col pred_prob --target-prob 0.5 --side {{side}}")
    changed = cfs[feats].astype(str).ne(query[feats].astype(str).iloc[0]).sum(axis=1)
    print(cfs.assign(n_changed=changed.values).to_string(index=False))
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a DiCE counterfactual script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", required=True)
    ap.add_argument("--train", required=True)
    ap.add_argument("--rows", required=True)
    ap.add_argument("--target", required=True)
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--constraints", help="JSON with a 'features' block (mutable, direction, min, max, integer)")
    ap.add_argument("--row", type=int, nargs="+", default=[0])
    ap.add_argument("--method", choices=["random", "genetic", "kdtree"], default="random")
    ap.add_argument("--total", type=int, default=4, help="counterfactuals per row")
    ap.add_argument("--desired-class", default="opposite", help='"opposite" or a class index')
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.total < 1 or min(a.row) < 0:
        print("error: --total must be >= 1 and rows >= 0", file=sys.stderr)
        return 2
    desired = a.desired_class if a.desired_class == "opposite" else int(a.desired_class)
    text = TEMPLATE.format(model=a.model, train=a.train, rows=a.rows, target=a.target, drop=a.drop,
                           constraints=a.constraints, row_ids=a.row, method=a.method, total=a.total,
                           desired=desired, script=os.path.basename(a.output or "dice_run.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {a.method}, {a.total} counterfactual(s) for rows {a.row}", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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