#!/usr/bin/env python3
"""Generate a production-shaped scikit-learn training script for a CSV: Pipeline + ColumnTransformer,
the right CV splitter, optional hyperparameter search, evaluation, persistence and a predict() helper.

Column roles are inferred from a sample (numeric, low-/high-cardinality categorical, datetime, text, ID,
constant) and can be overridden. Generated pipeline per model family:
  linear   numeric: median impute (+indicator) -> StandardScaler; categorical: OneHot(min_frequency);
           high-card: TargetEncoder; text: TF-IDF; dates: DateFeatures transformer -> scaled
  hgb      HistGradientBoosting with native categoricals (OrdinalEncoder, categorical_features mask),
           no scaling; high-card: TargetEncoder
  rf       RandomForest: ordinal-encoded categoricals, median imputation
CV: stratified K-fold (classification), K-fold (regression), GroupKFold (--group-col) or
TimeSeriesSplit on rows sorted by --time-col (with --gap). Search: none | random | halving.

The generated script needs pandas, scikit-learn >= 1.4 and joblib. This generator is stdlib only.

Usage:
    python gen_pipeline.py data.csv --target churned --model hgb --time-col signup_date -o train.py
    python gen_pipeline.py data.csv --target price --model linear --search random --drop house_id
    python gen_pipeline.py data.csv --target churned --group-col customer_id --text-cols notes

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

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

MISSING = {"", "na", "n/a", "nan", "null", "none", "?"}
DATE_RE = re.compile(r"^\d{4}[-/]\d{1,2}[-/]\d{1,2}")


def is_num(v):
    try:
        return math.isfinite(float(v))
    except ValueError:
        return False


def roles(header, rows, target, drop, time_col, group_col, text_cols, cat_cols):
    out = {"numeric": [], "categorical": [], "high_card": [], "date": [], "text": list(text_cols), "dropped": []}
    for c in header:
        if c in (target, group_col) or c in drop or c in text_cols:
            continue
        vals = [(r.get(c) or "").strip() for r in rows]
        present = [v for v in vals if v.lower() not in MISSING]
        if not present or len(set(present)) == 1:
            out["dropped"].append((c, "empty or constant"))
            continue
        k = len(set(present))
        if c in cat_cols:
            out["high_card" if k > 50 else "categorical"].append(c)
        elif c == time_col or sum(bool(DATE_RE.match(v)) for v in present) / len(present) > 0.9:
            out["date"].append(c)
        elif all(is_num(v) for v in present):
            if re.search(r"(^id$|_id$)", c, re.I) and k > 0.9 * len(present):
                out["dropped"].append((c, "numeric ID"))
            elif re.search(r"(zip|postcode|_code$)", c, re.I):
                out["high_card" if k > 50 else "categorical"].append(c)
            else:
                out["numeric"].append(c)
        elif k > 0.9 * len(present):
            mean_len = sum(len(v) for v in present) / len(present)
            if mean_len > 30:
                out["text"].append(c)
            else:
                out["dropped"].append((c, "near-unique text (ID)"))
        elif k > 50:
            out["high_card"].append(c)
        else:
            out["categorical"].append(c)
    return out


HEADER = '''"""Train {task} model for {target!r}. Generated by gen_pipeline.py - review, then version this file.

Run:   python {script} [--data {data}]
Needs: pandas, scikit-learn >= 1.4, joblib
"""
import argparse
import json

import joblib
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin, clone
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.model_selection import {splitter}, cross_validate{search_import}
from sklearn.pipeline import Pipeline, make_pipeline
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler, TargetEncoder
{extra_imports}
TARGET = {target!r}
NUMERIC = {numeric!r}
CATEGORICAL = {categorical!r}
HIGH_CARD = {high_card!r}
DATES = {date!r}
TEXT = {text!r}
DROP = {dropped!r}          # inferred as unusable; never feed IDs to the model
{group_line}{time_line}SEED = 0


class DateFeatures(BaseEstimator, TransformerMixin):
    """Expand datetime columns into numeric parts. Stateless, so safe in any CV scheme."""

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        out = []
        for c in X.columns:
            d = pd.to_datetime(X[c], errors="coerce")
            out += [d.dt.year, d.dt.month, d.dt.day_of_week, d.dt.day]
        return np.column_stack(out).astype(float) if out else np.empty((len(X), 0))

    def get_feature_names_out(self, input_features=None):
        return np.array([f"{{c}}_{{p}}" for c in input_features for p in ("year", "month", "dow", "day")])
'''

PREP_LINEAR = '''

def build_preprocessor():
    steps = []
    if NUMERIC:
        steps.append(("num", make_pipeline(SimpleImputer(strategy="median", add_indicator=True), StandardScaler()), NUMERIC))
    if CATEGORICAL:
        steps.append(("cat", make_pipeline(SimpleImputer(strategy="most_frequent"),
                                           OneHotEncoder(handle_unknown="infrequent_if_exist", min_frequency=20)),
                      CATEGORICAL))
    if HIGH_CARD:
        steps.append(("hc", TargetEncoder(smooth="auto"), HIGH_CARD))          # cross-fitted in fit_transform
    if DATES:
        steps.append(("date", make_pipeline(DateFeatures(), SimpleImputer(strategy="median"), StandardScaler()), DATES))
    for col in TEXT:
        steps.append((f"txt_{{col}}", TfidfVectorizer(ngram_range=(1, 2), min_df=5, max_features=20000), col))
    return ColumnTransformer(steps, remainder="drop", sparse_threshold=0.3)
'''

TEXT_SVD = '''

class TextSVD(BaseEstimator, TransformerMixin):
    """TruncatedSVD that caps n_components at vocabulary size - 1 (small corpora, small CV folds)."""

    def __init__(self, n_components=50, random_state=0):
        self.n_components, self.random_state = n_components, random_state

    def fit(self, X, y=None):
        k = max(1, min(self.n_components, X.shape[1] - 1))
        self.svd_ = TruncatedSVD(n_components=k, random_state=self.random_state).fit(X)
        return self

    def transform(self, X):
        return self.svd_.transform(X)
'''

PREP_TREE = '''

def build_preprocessor():
    steps = []
    if CATEGORICAL:
        steps.append(("cat", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1,
                                            encoded_missing_value=-1, min_frequency=20), CATEGORICAL))
    if NUMERIC:
        steps.append(("num", {num_step}, NUMERIC))
    if HIGH_CARD:
        steps.append(("hc", TargetEncoder(smooth="auto"), HIGH_CARD))
    if DATES:
        steps.append(("date", DateFeatures(), DATES))
    for col in TEXT:
        steps.append((f"txt_{{col}}", make_pipeline(TfidfVectorizer(min_df=5, max_features=5000),
                                                     TextSVD(n_components=50, random_state=SEED)), col))
    return ColumnTransformer(steps, remainder="drop", sparse_threshold=0.0)
'''

MODELS = {
    ("linear", "classification"): ("LogisticRegression(C=1.0, max_iter=5000, class_weight=None)",
                                   "from sklearn.linear_model import LogisticRegression",
                                   '{"model__C": loguniform(1e-3, 1e2)}'),
    ("linear", "regression"): ("Ridge(alpha=1.0)", "from sklearn.linear_model import Ridge",
                               '{"model__alpha": loguniform(1e-3, 1e3)}'),
    ("hgb", "classification"): ("HistGradientBoostingClassifier(categorical_features=CAT_MASK, early_stopping=True, "
                                "random_state=SEED)", "from sklearn.ensemble import HistGradientBoostingClassifier",
                                '{"model__learning_rate": loguniform(0.02, 0.3), "model__max_leaf_nodes": randint(8, 64), '
                                '"model__min_samples_leaf": randint(10, 200), "model__l2_regularization": loguniform(1e-4, 10)}'),
    ("hgb", "regression"): ("HistGradientBoostingRegressor(categorical_features=CAT_MASK, early_stopping=True, "
                            "random_state=SEED)", "from sklearn.ensemble import HistGradientBoostingRegressor",
                            '{"model__learning_rate": loguniform(0.02, 0.3), "model__max_leaf_nodes": randint(8, 64), '
                            '"model__min_samples_leaf": randint(10, 200), "model__l2_regularization": loguniform(1e-4, 10)}'),
    ("rf", "classification"): ("RandomForestClassifier(n_estimators=500, min_samples_leaf=5, n_jobs=-1, random_state=SEED)",
                               "from sklearn.ensemble import RandomForestClassifier",
                               '{"model__max_features": uniform(0.1, 0.8), "model__min_samples_leaf": randint(1, 50)}'),
    ("rf", "regression"): ("RandomForestRegressor(n_estimators=500, min_samples_leaf=5, n_jobs=-1, random_state=SEED)",
                           "from sklearn.ensemble import RandomForestRegressor",
                           '{"model__max_features": uniform(0.1, 0.8), "model__min_samples_leaf": randint(1, 50)}'),
}

BODY = '''

def build_model():
    return Pipeline([("prep", build_preprocessor()), ("model", {model})])


def load(path):
    df = pd.read_csv(path)
    df = df[df[TARGET].notna()].copy()
{load_extra}    return df


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default={data!r})
    ap.add_argument("--out", default="model.joblib")
    args = ap.parse_args()
    df = load(args.data)
{target_prep}    X = df.drop(columns=[TARGET])
    cv = {cv_expr}
    model = build_model()
{search_block}    res = cross_validate(model, X, y, cv=cv, scoring=SCORING, return_train_score=True, n_jobs=1{groups_kw})
    report = {{m: {{"test_mean": float(np.mean(res[f"test_{{m}}"])), "test_std": float(np.std(res[f"test_{{m}}"])),
                   "train_mean": float(np.mean(res[f"train_{{m}}"]))}} for m in SCORING}}
    print(json.dumps(report, indent=2))
    final = clone(model).fit(X, y)
    joblib.dump({{"model": final, "target": TARGET, "columns": list(X.columns){classes_entry}}}, args.out)
    print(f"saved {{args.out}}")


def predict(path_or_df, model_path="model.joblib"):
    """Score new rows with the saved pipeline (same preprocessing as training)."""
    bundle = joblib.load(model_path)
    X = pd.read_csv(path_or_df) if isinstance(path_or_df, str) else path_or_df
    X = X.reindex(columns=bundle["columns"])
    return {predict_expr}


if __name__ == "__main__":
    main()
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a scikit-learn Pipeline training script for a CSV.",
                                 epilog=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv")
    ap.add_argument("--target", required=True)
    ap.add_argument("--task", choices=["auto", "classification", "regression"], default="auto")
    ap.add_argument("--model", choices=["linear", "hgb", "rf"], default="hgb")
    ap.add_argument("--time-col")
    ap.add_argument("--gap", type=int, default=0, help="TimeSeriesSplit gap (rows) between train and test folds")
    ap.add_argument("--group-col")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--search", choices=["none", "random", "halving"], default="none")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--text-cols", nargs="*", default=[])
    ap.add_argument("--cat-cols", nargs="*", default=[], help="force these columns to be categorical")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    try:
        with open(a.csv, newline="", encoding="utf-8-sig") as fh:
            reader = csv.DictReader(fh)
            header = reader.fieldnames or []
            rows = list(islice(reader, 5000))
    except OSError as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    for c in [a.target, a.time_col, a.group_col, *a.text_cols, *a.cat_cols]:
        if c and c not in header:
            print(f"error: column {c!r} not found", file=sys.stderr)
            return 2
    tv = {(r.get(a.target) or "").strip() for r in rows} - {""}
    task = a.task if a.task != "auto" else ("classification" if len(tv) <= 20 or not all(is_num(v) for v in tv)
                                            else "regression")
    rl = roles(header, rows, a.target, set(a.drop), a.time_col, a.group_col, set(a.text_cols), set(a.cat_cols))
    if a.time_col and a.time_col in rl["date"]:
        rl["date"].remove(a.time_col)        # ordering only; raw calendar position of the split column is not a feature
        rl["dropped"].append((a.time_col, "time column used for ordering"))
    if a.group_col:
        rl["dropped"].append((a.group_col, "group key for CV"))
    model, model_import, dist = MODELS[(a.model, task)]
    # leaf-size ranges must fit the smallest training fold (TimeSeriesSplit's first fold is ~1/(k+1) of the rows)
    leaf_hi = max(10, min(200, len(rows) // (4 * (a.folds + 1))))
    dist = dist.replace("randint(10, 200)", f"randint(5, {leaf_hi})").replace("randint(1, 50)", f"randint(1, {max(5, leaf_hi // 2)})")
    extra = [model_import]
    if rl["text"]:
        extra.append("from sklearn.feature_extraction.text import TfidfVectorizer")
        if a.model != "linear":
            extra.append("from sklearn.decomposition import TruncatedSVD")
    if a.search != "none":
        extra.append("from scipy.stats import loguniform, randint, uniform")
    if a.search == "halving":
        extra.append("from sklearn.experimental import enable_halving_search_cv  # noqa: F401  (must precede the import)")
        extra.append("from sklearn.model_selection import HalvingRandomSearchCV")
    binary = task == "classification" and len(tv) == 2
    if task == "classification":
        scoring = '["roc_auc", "average_precision", "neg_log_loss"]' if binary else '["accuracy", "f1_macro", "neg_log_loss"]'
        primary = "average_precision" if binary else "f1_macro"
    else:
        scoring = '["neg_mean_absolute_error", "neg_root_mean_squared_error", "r2"]'
        primary = "neg_mean_absolute_error"
    if a.time_col:
        splitter, cv_expr = "TimeSeriesSplit", f"TimeSeriesSplit(n_splits={a.folds}, gap={a.gap})   # rows sorted by time in load()"
        load_extra = f"    df = df.sort_values({a.time_col!r}).reset_index(drop=True)\n"
    elif a.group_col:
        splitter = "StratifiedGroupKFold" if task == "classification" else "GroupKFold"
        cv_expr = (f"StratifiedGroupKFold(n_splits={a.folds}, shuffle=True, random_state=SEED)" if task == "classification"
                   else f"GroupKFold(n_splits={a.folds})")
        load_extra = ""
    else:
        splitter = "StratifiedKFold" if task == "classification" else "KFold"
        cv_expr = f"{splitter}(n_splits={a.folds}, shuffle=True, random_state=SEED)"
        load_extra = ""
    groups_kw = ", groups=groups" if a.group_col else ""
    if task == "classification" and binary:
        target_prep = ("    positive = df[TARGET].value_counts().idxmin()        # minority class = positive\n"
                       "    y = (df[TARGET] == positive).astype(int)\n")
        classes_entry, predict_expr = ', "positive": positive', 'bundle["model"].predict_proba(X)[:, 1]'
    elif task == "classification":
        target_prep = "    y = df[TARGET].astype(str)\n"
        classes_entry, predict_expr = "", 'bundle["model"].predict_proba(X)'
    else:
        target_prep = "    y = df[TARGET].astype(float)\n"
        classes_entry, predict_expr = "", 'bundle["model"].predict(X)'
    if a.group_col:
        target_prep += f"    groups = df.pop({a.group_col!r})\n"
    target_prep += f"    df = df.drop(columns=[c for c, _ in DROP if c in df.columns and c != TARGET])\n"
    search_block = ""
    if a.search != "none":
        cls = "RandomizedSearchCV" if a.search == "random" else "HalvingRandomSearchCV"
        n_iter = "n_iter=40, " if a.search == "random" else "factor=3, "
        search_block = (f"    search = {cls}(model, {dist}, {n_iter}scoring={primary!r}, cv=cv, random_state=SEED, n_jobs=-1)\n"
                        f"    search.fit(X, y{groups_kw})\n"
                        "    print(\"best params:\", search.best_params_, \"inner score:\", round(search.best_score_, 4))\n"
                        "    model = search.best_estimator_\n"
                        "    # cross_validate below re-scores the tuned configuration on the same folds (optimistic);\n"
                        "    # for an unbiased estimate wrap the search itself in cross_validate (nested CV).\n")
    cat_mask = ""
    if a.model == "hgb":
        cat_mask = ("\n# HistGradientBoosting: the ordinal-encoded CATEGORICAL block is first in the ColumnTransformer, so\n"
                    "# its output columns are 0..len(CATEGORICAL)-1 whatever the width of the later blocks\n"
                    "CAT_MASK = list(range(len(CATEGORICAL))) or None\n")
    num_step = '"passthrough"' if a.model == "hgb" else 'SimpleImputer(strategy="median")'
    script_name = os.path.basename(a.output or "train.py")
    search_import = ""
    if a.search == "random":
        search_import = ", RandomizedSearchCV"
    text = HEADER.format(task=task, target=a.target, script=script_name, data=a.csv, splitter=splitter,
                         search_import=search_import, extra_imports="\n".join(extra) + "\n",
                         numeric=rl["numeric"], categorical=rl["categorical"], high_card=rl["high_card"],
                         date=rl["date"], text=rl["text"], dropped=rl["dropped"],
                         group_line=f"GROUP = {a.group_col!r}\n" if a.group_col else "",
                         time_line=f"TIME = {a.time_col!r}\n" if a.time_col else "")
    text += f"SCORING = {scoring}\n" + cat_mask
    if rl["text"] and a.model != "linear":
        text += TEXT_SVD
    text += (PREP_LINEAR if a.model == "linear" else PREP_TREE.format(num_step=num_step))
    text += BODY.format(model=model, data=a.csv, load_extra=load_extra, target_prep=target_prep, cv_expr=cv_expr,
                        search_block=search_block, groups_kw=groups_kw, classes_entry=classes_entry,
                        predict_expr=predict_expr)
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: task={task}, model={a.model}, cv={splitter}, numeric={len(rl['numeric'])}, "
              f"categorical={len(rl['categorical'])}, high_card={len(rl['high_card'])}, dates={len(rl['date'])}, "
              f"text={len(rl['text'])}, dropped={len(rl['dropped'])}", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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