#!/usr/bin/env python3
"""Generate a runnable scikit-learn baseline script tailored to a CSV's columns.

Reads the header and a sample of rows, sorts columns into numeric / categorical / datetime /
dropped (IDs, near-unique text), and writes a script that:
  - splits honestly: time-ordered holdout (--time-col), grouped (--group-col) or stratified random
  - scores DummyClassifier/DummyRegressor, a linear model (LogisticRegression / Ridge) and
    HistGradientBoosting inside one Pipeline + ColumnTransformer each (no preprocessing leakage)
  - prints a comparison table with the metric family for the task

The generated script needs pandas and scikit-learn >= 1.2. This generator itself is stdlib only.

Usage:
    python scaffold_baseline.py data.csv --target churned --time-col signup_date > baseline.py
    python scaffold_baseline.py data.csv --target price --drop house_id -o baseline.py
    python scaffold_baseline.py data.csv --target churned --group-col customer_id --task classification

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", "?"}


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


def looks_date(v):
    v = v.strip()
    return len(v) >= 8 and v[:4].isdigit() and v[4] in "-/" and v[5:7].isdigit()


def classify(header, rows, target, drop, time_col, group_col):
    numeric, categorical, dates, dropped = [], [], [], []
    n = len(rows)
    for c in header:
        if c in (target, time_col, group_col) or c in drop:
            continue
        vals = [r.get(c, "") for r in rows if (r.get(c) or "").strip().lower() not in MISSING]
        if not vals:
            dropped.append((c, "empty"))
            continue
        distinct = len(set(vals))
        if all(is_num(v) for v in vals):
            if distinct == n and all(float(v).is_integer() for v in vals) and c.lower().endswith("id"):
                dropped.append((c, "integer ID"))
            elif c.lower().endswith(("zip", "zip_code", "postcode", "code")) and distinct > 50:
                categorical.append(c)   # numeric-looking codes are categories, not quantities
            else:
                numeric.append(c)
        elif sum(looks_date(v) for v in vals) / len(vals) > 0.9:
            dates.append(c)
        elif distinct > 0.9 * len(vals):
            dropped.append((c, "near-unique text (ID or free text)"))
        else:
            categorical.append(c)
    return numeric, categorical, dates, dropped


TEMPLATE = '''"""Baseline models for {target!r} ({task}). Generated by scaffold_baseline.py; edit freely.

Run: python {script_name}   (needs pandas, scikit-learn >= 1.2)
"""
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.dummy import {dummy}
from sklearn.ensemble import {hgb}
from sklearn.impute import SimpleImputer
from sklearn.linear_model import {linear}
from sklearn.metrics import {metric_imports}
from sklearn.model_selection import {split_import}
from sklearn.pipeline import Pipeline, make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

DATA = {data_path!r}
TARGET = {target!r}
NUMERIC = {numeric!r}
CATEGORICAL = {categorical!r}
DATES = {dates!r}          # expanded into numeric parts below
# Dropped automatically: {dropped}

df = pd.read_csv(DATA)
df = df[df[TARGET].notna()].copy()
for col in DATES:
    d = pd.to_datetime(df[col], errors="coerce")
    df[col + "_year"], df[col + "_month"], df[col + "_dow"] = d.dt.year, d.dt.month, d.dt.day_of_week
    NUMERIC = NUMERIC + [col + "_year", col + "_month", col + "_dow"]
{target_prep}
{split_code}

def preprocess(scale):
    num_steps = [SimpleImputer(strategy="median", add_indicator=True)] + ([StandardScaler()] if scale else [])
    num = make_pipeline(*num_steps)
    cat = make_pipeline(SimpleImputer(strategy="most_frequent"),
                        OneHotEncoder(handle_unknown="infrequent_if_exist", min_frequency=10, sparse_output=False))
    return ColumnTransformer([("num", num, NUMERIC), ("cat", cat, CATEGORICAL)], remainder="drop")


models = {{
    "dummy": Pipeline([("prep", preprocess(False)), ("model", {dummy}({dummy_args}))]),
    "linear": Pipeline([("prep", preprocess(True)), ("model", {linear}({linear_args}))]),
    "hist_gbm": Pipeline([("prep", preprocess(False)), ("model", {hgb}(random_state=0))]),
}}

rows = []
for name, model in models.items():
    model.fit(X_train, y_train)
{score_code}
report = pd.DataFrame(rows).set_index("model")
print(report.round(4).to_string())
{leak_check}print("\\nShip only what beats 'dummy' and today's rule by more than the noise between repeated splits.")
'''

CLS = dict(dummy="DummyClassifier", hgb="HistGradientBoostingClassifier", linear="LogisticRegression",
           metric_imports="average_precision_score, balanced_accuracy_score, brier_score_loss, log_loss, roc_auc_score",
           dummy_args='strategy="prior"', linear_args="max_iter=2000, C=1.0",
           leak_check='if report["roc_auc"].max() > 0.97:\n'
                      '    print("\\nWARNING: ROC-AUC > 0.97 on a first baseline usually means target leakage; "\n'
                      '          "run target-leakage-detector before trusting it.")\n',
           target_prep='POSITIVE = df[TARGET].value_counts().idxmin()   # minority class as positive\n'
                       'y_all = (df[TARGET] == POSITIVE).astype(int)\n',
           score_code='''    p = model.predict_proba(X_test)[:, 1]
    rows.append({"model": name, "roc_auc": roc_auc_score(y_test, p), "pr_auc": average_precision_score(y_test, p),
                 "log_loss": log_loss(y_test, p), "brier": brier_score_loss(y_test, p),
                 "balanced_acc@prior": balanced_accuracy_score(y_test, p >= y_train.mean())})''')
REG = dict(dummy="DummyRegressor", hgb="HistGradientBoostingRegressor", linear="Ridge",
           metric_imports="mean_absolute_error, mean_squared_error, r2_score",
           dummy_args='strategy="median"', linear_args="alpha=1.0",
           leak_check='if report["r2"].max() > 0.98:\n'
                      '    print("\\nWARNING: R2 > 0.98 on a first baseline usually means target leakage.")\n',
           target_prep="y_all = df[TARGET].astype(float)\n",
           score_code='''    p = model.predict(X_test)
    rows.append({"model": name, "mae": mean_absolute_error(y_test, p),
                 "rmse": float(np.sqrt(mean_squared_error(y_test, p))), "r2": r2_score(y_test, p),
                 "wape": float(np.abs(y_test - p).sum() / np.abs(y_test).sum())})''')


def split_code(time_col, group_col, task, test_frac):
    if time_col:
        return ("split_import", "train_test_split", f'''df = df.sort_values({time_col!r})
y_all = y_all.loc[df.index]
cut = int(len(df) * {1 - test_frac:g})          # latest {test_frac:.0%} of rows = test set
X_train, X_test = df.iloc[:cut], df.iloc[cut:]
y_train, y_test = y_all.iloc[:cut], y_all.iloc[cut:]
''')
    if group_col:
        return ("split_import", "GroupShuffleSplit", f'''gss = GroupShuffleSplit(n_splits=1, test_size={test_frac:g}, random_state=0)
train_idx, test_idx = next(gss.split(df, y_all, groups=df[{group_col!r}]))   # no {group_col} in both sets
X_train, X_test = df.iloc[train_idx], df.iloc[test_idx]
y_train, y_test = y_all.iloc[train_idx], y_all.iloc[test_idx]
''')
    strat = "stratify=y_all, " if task == "classification" else ""
    return ("split_import", "train_test_split", f'''X_train, X_test, y_train, y_test = train_test_split(
    df, y_all, test_size={test_frac:g}, {strat}random_state=0)
''')


def main():
    ap = argparse.ArgumentParser(description="Generate a scikit-learn baseline 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("--time-col")
    ap.add_argument("--group-col")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--test-frac", type=float, default=0.25)
    ap.add_argument("-o", "--output", help="write here instead of stdout")
    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
    if a.target not in header:
        print(f"error: target {a.target!r} not in columns", file=sys.stderr)
        return 2
    tvals = {(r.get(a.target) or "").strip() for r in rows} - {""}
    task = a.task
    if task == "auto":
        task = "classification" if len(tvals) <= 20 or not all(is_num(v) for v in tvals) else "regression"
    if task == "classification" and len(tvals) != 2:
        print(f"note: {len(tvals)} classes; the template scores the minority class one-vs-rest. For multiclass, "
              "swap metrics for macro-F1 and log_loss on predict_proba.", file=sys.stderr)
    numeric, categorical, dates, dropped = classify(header, rows, a.target, set(a.drop), a.time_col, a.group_col)
    _, split_name, code = split_code(a.time_col, a.group_col, task, a.test_frac)
    parts = dict(CLS if task == "classification" else REG)
    script = TEMPLATE.format(target=a.target, task=task, script_name=os.path.basename(a.output or "baseline.py"), data_path=a.csv,
                             numeric=numeric, categorical=categorical, dates=dates,
                             dropped=", ".join(f"{c} ({why})" for c, why in dropped) or "none",
                             split_import=split_name, split_code=code, **parts)
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(script)
        print(f"wrote {a.output}: {len(numeric)} numeric, {len(categorical)} categorical, {len(dates)} date, "
              f"{len(dropped)} dropped", file=sys.stderr)
    else:
        sys.stdout.write(script)
    return 0


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