#!/usr/bin/env python3
"""Generate an uplift modelling script: S-, T- and X-learners and the class-transformation approach.

The script it writes, for randomised data (treatment 1/0, outcome, covariates):
  - one-hot encodes categorical covariates and splits train/holdout stratified on treatment x outcome
  - fits on the training part, with LightGBM (or --base logreg) as the base learner:
      s_learner   one model y ~ X + t; uplift = f(X, 1) - f(X, 0)
      t_learner   separate models for treated and control; uplift = f1(X) - f0(X)
      x_learner   T-learner first stage, imputed effects D1 = y - f0(X) (treated), D0 = f1(X) - y (control),
                  second-stage regressors, combined with the propensity e(X): e * g0 + (1 - e) * g1
      class_transform  (Jaszewska & Lambert): z = 1 if (t=1, y=1) or (t=0, y=0); uplift = 2 P(z=1|X) - 1,
                  valid for binary outcomes and a 50/50 assignment (rows are reweighted otherwise)
  - writes holdout_scores.csv (treatment, outcome, one uplift column per learner, plus the covariates) for
    uplift_eval.py, which reports Qini, AUUC, deciles and targeting value on identical holdout rows
  - prints a quick Qini-free summary (mean predicted uplift vs observed ATE on the holdout) and saves the models
Needs pandas, numpy, scikit-learn, lightgbm (unless --base logreg), joblib.

Usage:
    python gen_uplift_script.py --data uplift.csv --treatment-col treatment --outcome-col converted --covariates age segment -o uplift.py
    python gen_uplift_script.py --data exp.csv --covariates age plan region --learners t_learner x_learner --base logreg

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

import argparse
import os
import sys

TEMPLATE = '''"""Uplift models for {data!r}. Generated by gen_uplift_script.py.

Run:  python {script}
Then: python uplift_eval.py holdout_scores.csv --treatment-col {treatment} --outcome-col {outcome} --scores {learners_space}
"""
import joblib
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.model_selection import train_test_split

DATA, T, Y, COVARIATES, LEARNERS, BASE, HOLDOUT, SEED = {data!r}, {treatment!r}, {outcome!r}, {covariates!r}, {learners!r}, {base!r}, {holdout}, 0

df = pd.read_csv(DATA).dropna(subset=[T, Y] + COVARIATES)
X_all = pd.get_dummies(df[COVARIATES], drop_first=False, dtype=float)
feats = list(X_all.columns)
data = pd.concat([X_all, df[[T, Y]]], axis=1)
binary = set(data[Y].unique()) <= {{0, 1}}
strat = data[T].astype(str) + "_" + (data[Y].astype(str) if binary else "")
train, hold = train_test_split(data, test_size=HOLDOUT, stratify=strat, random_state=SEED)
Xtr, ttr, ytr = train[feats], train[T].to_numpy(), train[Y].to_numpy()
Xho = hold[feats]


def clf():
    if BASE == "logreg":
        return LogisticRegression(max_iter=5000, C=1.0)
    from lightgbm import LGBMClassifier
    return LGBMClassifier(n_estimators=300, learning_rate=0.03, num_leaves=7, min_child_samples=50,
                          subsample=0.8, subsample_freq=1, random_state=SEED, verbose=-1)


def reg():
    if BASE == "logreg":
        return LinearRegression()
    from lightgbm import LGBMRegressor
    return LGBMRegressor(n_estimators=300, learning_rate=0.03, num_leaves=7, min_child_samples=50,
                         subsample=0.8, subsample_freq=1, random_state=SEED, verbose=-1)


def outcome_model():
    return clf() if binary else reg()


def pred(m, X):
    return m.predict_proba(X)[:, 1] if hasattr(m, "predict_proba") else m.predict(X)


models, scores = {{}}, {{}}
if "s_learner" in LEARNERS:
    m = outcome_model().fit(Xtr.assign(_t=ttr), ytr)
    scores["s_learner"] = pred(m, Xho.assign(_t=1)) - pred(m, Xho.assign(_t=0))
    models["s_learner"] = m
if "t_learner" in LEARNERS or "x_learner" in LEARNERS:
    m1 = outcome_model().fit(Xtr[ttr == 1], ytr[ttr == 1])
    m0 = outcome_model().fit(Xtr[ttr == 0], ytr[ttr == 0])
    if "t_learner" in LEARNERS:
        scores["t_learner"] = pred(m1, Xho) - pred(m0, Xho)
        models["t_learner"] = (m1, m0)
if "x_learner" in LEARNERS:
    d1 = ytr[ttr == 1] - pred(m0, Xtr[ttr == 1])
    d0 = pred(m1, Xtr[ttr == 0]) - ytr[ttr == 0]
    g1 = reg().fit(Xtr[ttr == 1], d1)
    g0 = reg().fit(Xtr[ttr == 0], d0)
    prop = LogisticRegression(max_iter=5000).fit(Xtr, ttr)
    e = prop.predict_proba(Xho)[:, 1]
    scores["x_learner"] = e * g0.predict(Xho) + (1 - e) * g1.predict(Xho)
    models["x_learner"] = (g1, g0, prop)
if "class_transform" in LEARNERS:
    if not binary:
        print("class_transform skipped: needs a binary outcome")
    else:
        z = ((ttr == 1) & (ytr == 1)) | ((ttr == 0) & (ytr == 0))
        p = ttr.mean()
        w = np.where(ttr == 1, 0.5 / p, 0.5 / (1 - p))              # reweight to a 50/50 assignment
        m = clf().fit(Xtr, z.astype(int), sample_weight=w)
        scores["class_transform"] = 2 * pred(m, Xho) - 1
        models["class_transform"] = m

out = hold[[T, Y]].copy()
for k, v in scores.items():
    out[k] = v
out = pd.concat([out, df.loc[hold.index, COVARIATES]], axis=1)
out.to_csv("holdout_scores.csv", index=False)
ate = hold.loc[hold[T] == 1, Y].mean() - hold.loc[hold[T] == 0, Y].mean()
print(f"holdout n={{len(hold)}}, observed ATE {{ate:+.4f}}")
for k, v in scores.items():
    print(f"{{k:<16}} mean predicted uplift {{np.mean(v):+.4f}}   share predicted negative {{np.mean(np.asarray(v) < 0):.0%}}")
joblib.dump({{"models": models, "features": feats}}, "uplift_models.joblib")
print("wrote holdout_scores.csv and uplift_models.joblib")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate an uplift modelling script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--data", required=True)
    ap.add_argument("--treatment-col", default="treatment")
    ap.add_argument("--outcome-col", default="converted")
    ap.add_argument("--covariates", nargs="+", required=True)
    ap.add_argument("--learners", nargs="+", choices=["s_learner", "t_learner", "x_learner", "class_transform"],
                    default=["s_learner", "t_learner", "x_learner", "class_transform"])
    ap.add_argument("--base", choices=["lgbm", "logreg"], default="lgbm")
    ap.add_argument("--holdout", type=float, default=0.4)
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if not 0.2 <= a.holdout <= 0.6:
        print("error: --holdout must be between 0.2 and 0.6", file=sys.stderr)
        return 2
    text = TEMPLATE.format(data=a.data, treatment=a.treatment_col, outcome=a.outcome_col, covariates=a.covariates,
                           learners=a.learners, learners_space=" ".join(a.learners), base=a.base, holdout=a.holdout,
                           script=os.path.basename(a.output or "uplift.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {', '.join(a.learners)} with {a.base}", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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