#!/usr/bin/env python3
"""Generate a fair classifier bake-off script: every candidate on identical CV folds, same metrics,
per-fold results saved for statistical comparison.

Candidates (--models): logreg, hgb, rf, et, lgbm, xgb, catboost, ebm. Optional libraries (lightgbm,
xgboost, catboost, interpret) are imported lazily; missing ones are skipped with a message.
Each candidate gets the preprocessing its family needs (scaling + one-hot for logreg; ordinal codes
for sklearn trees; pandas category dtype for LightGBM/XGBoost; raw strings for CatBoost/EBM).

The generated script writes bakeoff_folds.csv (model, fold, metric columns, fit_seconds,
predict_ms_per_1k) - feed it to ml-model-comparison-reporter/scripts/compare_models.py - and
prints a summary table.

Usage:
    python bakeoff_scaffold.py data.csv --target churned --models logreg hgb lgbm catboost -o bakeoff.py
    python bakeoff_scaffold.py data.csv --target churned --group-col customer_id --drop customer_id
    python bakeoff_scaffold.py data.csv --target churned --time-col signup_date --folds 5

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

import argparse
import csv
import math
import os
import sys
from itertools import islice

MISSING = {"", "na", "n/a", "nan", "null", "none", "?"}
KNOWN = ["logreg", "hgb", "rf", "et", "lgbm", "xgb", "catboost", "ebm"]

TEMPLATE = '''"""Classifier bake-off for {target!r}. Generated by bakeoff_scaffold.py.

Every model sees the same folds; results per fold go to bakeoff_folds.csv for paired comparison.
Run: python {script}   (pandas, scikit-learn >= 1.4; optional: lightgbm, xgboost, catboost, interpret)
"""
import time

import numpy as np
import pandas as pd
from sklearn.base import clone
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import ExtraTreesClassifier, HistGradientBoostingClassifier, RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, brier_score_loss, log_loss, roc_auc_score
from sklearn.model_selection import {splitter}
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler

DATA, TARGET = {data!r}, {target!r}
DROP = {drop!r}
MODELS = {models!r}
SEED = 0

df = pd.read_csv(DATA)
df = df[df[TARGET].notna()]
{order_line}positive = df[TARGET].value_counts().idxmin()
y = (df[TARGET] == positive).astype(int).to_numpy()
{groups_line}X = df.drop(columns=[TARGET] + [c for c in DROP if c in df.columns])
CAT = [c for c in X.columns if not pd.api.types.is_numeric_dtype(X[c])]
NUM = [c for c in X.columns if c not in CAT]
folds = list({cv_expr})


def onehot_scaled():
    return ColumnTransformer([
        ("num", make_pipeline(SimpleImputer(strategy="median", add_indicator=True), StandardScaler()), NUM),
        ("cat", make_pipeline(SimpleImputer(strategy="most_frequent"),
                              OneHotEncoder(handle_unknown="infrequent_if_exist", min_frequency=10)), CAT)])


def ordinal(impute):
    num = SimpleImputer(strategy="median") if impute else "passthrough"
    return ColumnTransformer([
        ("cat", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1, encoded_missing_value=-1), CAT),
        ("num", num, NUM)])


def as_category(frame, cats):
    out = frame.copy()
    for c in CAT:
        out[c] = pd.Categorical(out[c], categories=cats[c])
    return out


def make(name):
    """Return (estimator, kind); kind says how to feed data."""
    if name == "logreg":
        return make_pipeline(onehot_scaled(), LogisticRegression(C=1.0, max_iter=5000)), "sk"
    if name == "hgb":
        return make_pipeline(ordinal(False), HistGradientBoostingClassifier(
            categorical_features=list(range(len(CAT))) or None, learning_rate=0.05, early_stopping=True,
            random_state=SEED)), "sk"
    if name == "rf":
        return make_pipeline(ordinal(True), RandomForestClassifier(n_estimators=500, min_samples_leaf=3, n_jobs=-1,
                                                                   random_state=SEED)), "sk"
    if name == "et":
        return make_pipeline(ordinal(True), ExtraTreesClassifier(n_estimators=500, min_samples_leaf=3, n_jobs=-1,
                                                                 random_state=SEED)), "sk"
    if name == "lgbm":
        import lightgbm as lgb
        return lgb.LGBMClassifier(n_estimators=400, learning_rate=0.05, num_leaves=31, min_child_samples=20,
                                  subsample=0.8, subsample_freq=1, colsample_bytree=0.8, random_state=SEED,
                                  verbose=-1), "category"
    if name == "xgb":
        import xgboost as xgb
        return xgb.XGBClassifier(n_estimators=400, learning_rate=0.05, max_depth=6, subsample=0.8,
                                 colsample_bytree=0.8, tree_method="hist", enable_categorical=True,
                                 random_state=SEED), "category"
    if name == "catboost":
        from catboost import CatBoostClassifier
        return CatBoostClassifier(iterations=600, learning_rate=0.05, depth=6, random_seed=SEED, verbose=0,
                                  cat_features=CAT), "raw"
    if name == "ebm":
        from interpret.glassbox import ExplainableBoostingClassifier
        return ExplainableBoostingClassifier(random_state=SEED), "raw"
    raise ValueError(name)


rows = []
for name in MODELS:
    try:
        proto, kind = make(name)
    except ImportError as e:
        print(f"skip {{name}}: {{e}}")
        continue
    for k, (tr, te) in enumerate(folds):
        Xtr, Xte = X.iloc[tr], X.iloc[te]
        if kind == "category":
            cats = {{c: pd.Index(Xtr[c].dropna().unique()) for c in CAT}}
            Xtr, Xte = as_category(Xtr, cats), as_category(Xte, cats)
        elif kind == "raw":
            Xtr, Xte = Xtr.copy(), Xte.copy()
            for c in CAT:
                Xtr[c], Xte[c] = Xtr[c].astype(str), Xte[c].astype(str)
        est = clone(proto) if kind == "sk" else make(name)[0]
        t0 = time.perf_counter()
        est.fit(Xtr, y[tr])
        fit_s = time.perf_counter() - t0
        t0 = time.perf_counter()
        p = est.predict_proba(Xte)[:, 1]
        pred_ms = (time.perf_counter() - t0) / max(1, len(te)) * 1e6
        rows.append({{"model": name, "fold": k, "roc_auc": roc_auc_score(y[te], p),
                     "average_precision": average_precision_score(y[te], p),
                     "log_loss": log_loss(y[te], np.clip(p, 1e-15, 1 - 1e-15), labels=[0, 1]),
                     "brier": brier_score_loss(y[te], p), "fit_seconds": fit_s, "predict_ms_per_1k": pred_ms}})
    print(f"done {{name}}")

res = pd.DataFrame(rows)
res.to_csv("bakeoff_folds.csv", index=False)
summary = res.groupby("model").agg(["mean", "std"]).drop(columns="fold", level=0)
print(summary.round(4).to_string())
print("\\nPer-fold results in bakeoff_folds.csv; compare with ml-model-comparison-reporter (paired tests).")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a fair classifier bake-off script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv")
    ap.add_argument("--target", required=True)
    ap.add_argument("--models", nargs="+", default=["logreg", "hgb", "rf", "lgbm", "catboost"])
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--group-col")
    ap.add_argument("--time-col")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    bad = [m for m in a.models if m not in KNOWN]
    if bad:
        print(f"error: unknown model(s) {bad}; choose from {KNOWN}", file=sys.stderr)
        return 2
    try:
        with open(a.csv, newline="", encoding="utf-8-sig") as fh:
            reader = csv.DictReader(fh)
            header = reader.fieldnames or []
            list(islice(reader, 1))
    except OSError as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    for c in (a.target, a.group_col, a.time_col):
        if c and c not in header:
            print(f"error: column {c!r} not found", file=sys.stderr)
            return 2
    drop = list(a.drop)
    if a.time_col:
        splitter = "TimeSeriesSplit"
        cv_expr = f"TimeSeriesSplit(n_splits={a.folds}).split(X)"
        order_line = f"df = df.sort_values({a.time_col!r}).reset_index(drop=True)\n"
        drop.append(a.time_col)
        groups_line = ""
    elif a.group_col:
        splitter = "StratifiedGroupKFold"
        cv_expr = f"StratifiedGroupKFold(n_splits={a.folds}, shuffle=True, random_state=SEED).split(X, y, groups)"
        order_line = ""
        groups_line = f"groups = df[{a.group_col!r}].to_numpy()\n"
        if a.group_col not in drop:
            drop.append(a.group_col)
    else:
        splitter = "StratifiedKFold"
        cv_expr = f"StratifiedKFold(n_splits={a.folds}, shuffle=True, random_state=SEED).split(X, y)"
        order_line = groups_line = ""
    text = TEMPLATE.format(target=a.target, script=os.path.basename(a.output or "bakeoff.py"), splitter=splitter,
                           data=a.csv, drop=drop, models=a.models, order_line=order_line, groups_line=groups_line,
                           cv_expr=cv_expr)
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output} ({len(a.models)} models, {splitter})", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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