"""templates — train.py and predict.py source code for codegen export."""

from __future__ import annotations

_TRAIN_PY = '''\
#!/usr/bin/env python3
"""Train a LightGBM model and fit a probability calibrator.

Usage:
    python train.py train_data.csv
    python train.py train_data.parquet --no-calibration

Steps:
    1. Fit feature pipeline (learn category mappings -> pipeline_state.json)
    2. Train LightGBM on full data (-> model.txt)
    3. Generate OOF scores via CV for calibration
    4. Fit calibrator on OOF scores (-> calibrator.json)

Generated by lizyml -- https://github.com/nbx-liz/LizyML
"""
from __future__ import annotations

import argparse
import json
import logging
import math
import sys
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import KFold, StratifiedKFold

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s  %(message)s", datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

ROOT = Path(__file__).parent
ARTIFACTS = ROOT / "artifacts"

with open(ROOT / "config.json") as _f:
    CFG = json.load(_f)


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Feature Pipeline
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def fit_pipeline(df: pd.DataFrame) -> dict:
    """Learn category mappings and save pipeline state."""
    expected = CFG["feature_names"]
    missing = sorted(set(expected) - set(df.columns))
    if missing:
        raise ValueError(f"Missing columns: {missing}")

    mappings: dict[str, dict[str, int]] = {}
    for col in CFG["categorical_features"]:
        cats = sorted(str(v) for v in df[col].dropna().unique())
        mappings[col] = {v: i for i, v in enumerate(cats)}
        log.info("    %s: %d categories", col, len(cats))

    state = {
        "feature_names": expected,
        "categorical_features": CFG["categorical_features"],
        "category_mappings": mappings,
    }
    ARTIFACTS.mkdir(parents=True, exist_ok=True)
    with open(ARTIFACTS / "pipeline_state.json", "w") as f:
        json.dump(state, f, indent=2, ensure_ascii=False)
    return state


def transform(df: pd.DataFrame, state: dict) -> pd.DataFrame:
    """Apply fitted pipeline to a DataFrame."""
    X = df[state["feature_names"]].copy()
    for col, mapping in state.get("category_mappings", {}).items():
        if col in X.columns:
            X[col] = X[col].astype(str).map(mapping)  # unseen -> NaN
    return X


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  LightGBM Training
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def train_lgbm(X: pd.DataFrame, y: pd.Series, cat_cols: list[str]) -> lgb.Booster:
    """Train LightGBM with optional early stopping via holdout."""
    ds_full = lgb.Dataset(X, label=y, categorical_feature=cat_cols or "auto")
    callbacks: list = [lgb.log_evaluation(period=200)]
    train_set = ds_full
    valid_sets, valid_names = [ds_full], ["train"]

    ratio = CFG.get("validation_ratio", 0)
    es_rounds = CFG.get("early_stopping_rounds")
    if ratio > 0 and es_rounds:
        n = len(y)
        rng = np.random.default_rng(CFG["seed"])
        idx = rng.permutation(n)
        n_val = max(1, int(n * ratio))

        train_set = lgb.Dataset(
            X.iloc[idx[n_val:]], label=y.iloc[idx[n_val:]],
            categorical_feature=cat_cols or "auto",
        )
        valid_ds = lgb.Dataset(
            X.iloc[idx[:n_val]], label=y.iloc[idx[:n_val]],
            reference=train_set,
        )
        valid_sets = [train_set, valid_ds]
        valid_names = ["train", "valid"]
        callbacks.insert(0, lgb.early_stopping(es_rounds, verbose=True))
        log.info("    holdout: %d train / %d valid", n - n_val, n_val)

    booster = lgb.train(
        CFG["lgbm_params"], train_set,
        num_boost_round=CFG["num_boost_round"],
        valid_sets=valid_sets, valid_names=valid_names,
        callbacks=callbacks,
    )
    booster.save_model(str(ARTIFACTS / "model.txt"))
    log.info("    saved model.txt (best_iteration=%d)", booster.best_iteration)
    return booster


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Calibration
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _sigmoid(x: np.ndarray) -> np.ndarray:
    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))


def _generate_oof(X: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Lightweight CV to produce OOF raw scores (logits)."""
    n_splits = CFG.get("calibration_n_splits", 5)
    task = CFG["_task"]
    seed = CFG["seed"]

    kf = (StratifiedKFold if task == "binary" else KFold)(
        n_splits=n_splits, shuffle=True, random_state=seed,
    )
    oof = np.full(len(y), np.nan)
    for i, (trn, val) in enumerate(kf.split(X, y)):
        log.info("    CV fold %d/%d", i + 1, n_splits)
        ds = lgb.Dataset(X[trn], label=y[trn])
        bst = lgb.train(CFG["lgbm_params"], ds,
                        num_boost_round=CFG["num_boost_round"],
                        callbacks=[lgb.log_evaluation(0)])
        oof[val] = bst.predict(X[val], raw_score=True)
    return oof


def _fit_platt(scores: np.ndarray, y: np.ndarray) -> dict:
    lr = LogisticRegression(C=1.0, solver="lbfgs", max_iter=200)
    lr.fit(scores.reshape(-1, 1), y)
    return {"method": "platt",
            "a": float(lr.coef_[0, 0]), "b": float(lr.intercept_[0])}


def _fit_beta(scores: np.ndarray, y: np.ndarray) -> dict:
    from scipy.optimize import minimize
    s = np.clip(_sigmoid(scores), 1e-10, 1 - 1e-10)
    yf = y.astype(float)
    ls, l1s = np.log(s), np.log(1 - s)

    def nll(p):
        prob = np.clip(_sigmoid(p[0] * ls + p[1] * l1s + p[2]), 1e-10, 1 - 1e-10)
        return float(-np.sum(yf * np.log(prob) + (1 - yf) * np.log(1 - prob)))

    r = minimize(nll, x0=[1, 1, 0], method="L-BFGS-B")
    return {"method": "beta",
            "a": float(r.x[0]), "b": float(r.x[1]), "c": float(r.x[2])}


def _fit_isotonic(scores: np.ndarray, y: np.ndarray) -> dict:
    n = len(scores)
    params = {
        "objective": "binary", "metric": "binary_logloss",
        "monotone_constraints": [1], "monotone_constraints_method": "advanced",
        "num_leaves": 7, "max_depth": 3, "learning_rate": 0.03,
        "lambda_l2": 5.0, "min_data_in_leaf": max(1, math.ceil(n * 0.01)),
        "verbose": -1, "seed": CFG["seed"],
    }
    rng = np.random.default_rng(CFG["seed"])
    n_val = max(1, int(n * 0.1))
    idx = rng.permutation(n)
    X_cal = scores.reshape(-1, 1)

    ds_t = lgb.Dataset(X_cal[idx[n_val:]], label=y[idx[n_val:]].astype(float))
    ds_v = lgb.Dataset(X_cal[idx[:n_val]], label=y[idx[:n_val]].astype(float),
                       reference=ds_t)
    bst = lgb.train(params, ds_t, num_boost_round=1000,
                    valid_sets=[ds_v], valid_names=["valid"],
                    callbacks=[lgb.early_stopping(100, verbose=False),
                               lgb.log_evaluation(0)])
    bst.save_model(str(ARTIFACTS / "calibrator_model.txt"))
    return {"method": "isotonic", "model_file": "calibrator_model.txt"}


_CAL_FITTERS = {"platt": _fit_platt, "beta": _fit_beta, "isotonic": _fit_isotonic}


def fit_calibrator(X: np.ndarray, y: np.ndarray) -> dict | None:
    """Generate OOF scores and fit calibrator. Returns params or None."""
    method = CFG.get("calibration_method")
    if not method or CFG["_task"] != "binary":
        return None

    log.info("[3/4] Generating OOF scores ...")
    oof = _generate_oof(X, y)

    log.info("[4/4] Fitting %s calibrator ...", method)
    if method not in _CAL_FITTERS:
        raise ValueError(
            f"Unknown calibration method: {method!r}. "
            f"Supported: {list(_CAL_FITTERS)}"
        )
    params = _CAL_FITTERS[method](oof, y)
    with open(ARTIFACTS / "calibrator.json", "w") as f:
        json.dump(params, f, indent=2)
    log.info("    saved calibrator.json")
    return params


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Main
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def train(df: pd.DataFrame, *, calibrate: bool = True) -> None:
    target = CFG["_target_col"]
    y = df[target]
    X_raw = df.drop(columns=[target])

    log.info("[1/4] Fitting feature pipeline ...")
    state = fit_pipeline(X_raw)
    X = transform(X_raw, state)

    log.info("[2/4] Training LightGBM ...")
    cat_cols = [c for c in CFG["categorical_features"] if c in X.columns]
    train_lgbm(X, y, cat_cols)

    if calibrate:
        fit_calibrator(X.values, y.values)
    else:
        log.info("[3/4] Calibration skipped")
        log.info("[4/4] --")

    log.info("Done.")


def main() -> None:
    p = argparse.ArgumentParser(description="Train LightGBM (LizyML codegen)")
    p.add_argument("data", help="CSV or Parquet file")
    p.add_argument("--no-calibration", action="store_true")
    args = p.parse_args()

    path = Path(args.data)
    df = pd.read_parquet(path) if path.suffix == ".parquet" else pd.read_csv(path)

    target = CFG["_target_col"]
    if target not in df.columns:
        log.error(\'Target "%s" not found. Columns: %s\', target, list(df.columns))
        sys.exit(1)

    train(df, calibrate=not args.no_calibration)


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

_PREDICT_PY = '''\
#!/usr/bin/env python3
"""Run inference with a trained model.

Usage:
    python predict.py test_data.csv
    python predict.py test_data.csv -o predictions.csv

Generated by lizyml -- https://github.com/nbx-liz/LizyML
"""
from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s  %(message)s", datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

ROOT = Path(__file__).parent
ARTIFACTS = ROOT / "artifacts"

with open(ROOT / "config.json") as _f:
    CFG = json.load(_f)


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Feature Transform (predict-time: no re-fitting)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _load_pipeline() -> dict:
    path = ARTIFACTS / "pipeline_state.json"
    if not path.exists():
        raise FileNotFoundError(f"{path} not found. Run train.py first.")
    with open(path) as f:
        return json.load(f)


def transform(df: pd.DataFrame) -> pd.DataFrame:
    """Select expected columns and apply categorical encoding."""
    state = _load_pipeline()
    expected = state["feature_names"]

    missing = sorted(set(expected) - set(df.columns))
    if missing:
        raise ValueError(
            f"Missing {len(missing)} column(s): {missing}. "
            f"Expected: {expected}"
        )

    extra = sorted(set(df.columns) - set(expected))
    if extra:
        log.warning("Ignoring %d extra column(s): %s", len(extra), extra)

    X = df[expected].copy()
    for col, mapping in state.get("category_mappings", {}).items():
        if col in X.columns:
            X[col] = X[col].astype(str).map(mapping)  # unseen -> NaN
    return X


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Calibration (apply only)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _sigmoid(x: np.ndarray) -> np.ndarray:
    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))


def _load_calibrator() -> dict | None:
    path = ARTIFACTS / "calibrator.json"
    if not path.exists():
        return None
    with open(path) as f:
        return json.load(f)


def calibrate(raw_scores: np.ndarray, cal: dict) -> np.ndarray:
    """Map raw logits -> calibrated probabilities."""
    m = cal["method"]
    if m == "platt":
        return 1 / (1 + np.exp(-(cal["a"] * raw_scores + cal["b"])))
    if m == "beta":
        s = np.clip(_sigmoid(raw_scores), 1e-10, 1 - 1e-10)
        logit = cal["a"] * np.log(s) + cal["b"] * np.log(1 - s) + cal["c"]
        return np.clip(_sigmoid(logit), 0, 1)
    if m == "isotonic":
        model_file = Path(cal["model_file"]).name  # sanitize path
        bst = lgb.Booster(model_file=str(ARTIFACTS / model_file))
        return np.clip(bst.predict(raw_scores.reshape(-1, 1)), 0, 1)
    raise ValueError(
        f\'Unknown calibration method: "{m}". \'
        "Supported: platt, beta, isotonic"
    )


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Predict
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def predict(df: pd.DataFrame) -> dict[str, np.ndarray | None]:
    """Run inference. Returns {"pred": ..., "proba": ...}."""
    X = transform(df)
    booster = lgb.Booster(model_file=str(ARTIFACTS / "model.txt"))
    task = CFG["_task"]

    if task == "regression":
        return {"pred": np.asarray(booster.predict(X), dtype=np.float64),
                "proba": None}

    if task == "binary":
        proba = np.asarray(booster.predict(X), dtype=np.float64)
        cal = _load_calibrator()
        if cal:
            logits = np.asarray(booster.predict(X, raw_score=True), dtype=np.float64)
            proba = calibrate(logits, cal)
        return {"pred": (proba > 0.5).astype(np.int64), "proba": proba}

    if task == "multiclass":
        proba = np.asarray(booster.predict(X), dtype=np.float64)
        return {"pred": np.argmax(proba, axis=1).astype(np.int64), "proba": proba}

    raise ValueError(f\'Unknown task: "{task}"\')


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def main() -> None:
    p = argparse.ArgumentParser(description="Predict (LizyML codegen)")
    p.add_argument("data", help="CSV or Parquet file")
    p.add_argument("-o", "--output", default="predictions.csv")
    args = p.parse_args()

    path = Path(args.data)
    df = pd.read_parquet(path) if path.suffix == ".parquet" else pd.read_csv(path)
    result = predict(df)

    out = pd.DataFrame({"pred": result["pred"]})
    if result["proba"] is not None:
        if result["proba"].ndim == 1:
            out["proba"] = result["proba"]
        else:
            for i in range(result["proba"].shape[1]):
                out[f"proba_{i}"] = result["proba"][:, i]
    out.to_csv(args.output, index=False)
    log.info("Saved %d rows -> %s", len(out), args.output)


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

_TEST_EQUIVALENCE_PY = '''\
#!/usr/bin/env python3
"""Verify that codegen predictions match a LizyML reference.

Usage:
    python test_equivalence.py test_data.csv --reference reference_predictions.csv

Compares the output of predict.py against a reference CSV exported from
LizyML (``Model.predict()``).  The reference file must contain a ``pred``
column and optionally a ``proba`` column.

Generated by lizyml -- https://github.com/nbx-liz/LizyML
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s  %(message)s", datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

ROOT = Path(__file__).parent
ARTIFACTS = ROOT / "artifacts"

with open(ROOT / "config.json") as _f:
    CFG = json.load(_f)


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Prediction (inline — mirrors predict.py)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _load_pipeline() -> dict:
    path = ARTIFACTS / "pipeline_state.json"
    if not path.exists():
        raise FileNotFoundError(f"{path} not found. Run train.py first.")
    with open(path) as f:
        return json.load(f)


def _sigmoid(x: np.ndarray) -> np.ndarray:
    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))


def _transform(df: pd.DataFrame) -> np.ndarray:
    state = _load_pipeline()
    X = df[state["feature_names"]].copy()
    for col, mapping in state.get("category_mappings", {}).items():
        if col in X.columns:
            X[col] = X[col].astype(str).map(mapping)
    for col in X.columns:
        if not np.issubdtype(X[col].dtype, np.number):
            X[col] = pd.to_numeric(X[col], errors="coerce")
    return X.values


def _predict(df: pd.DataFrame) -> dict[str, np.ndarray | None]:
    X = _transform(df)
    booster = lgb.Booster(model_file=str(ARTIFACTS / "model.txt"))
    task = CFG["_task"]

    if task == "regression":
        return {"pred": np.asarray(booster.predict(X), dtype=np.float64),
                "proba": None}

    if task == "binary":
        cal_path = ARTIFACTS / "calibrator.json"
        if cal_path.exists():
            with open(cal_path) as f:
                cal = json.load(f)
            logits = np.asarray(booster.predict(X, raw_score=True), dtype=np.float64)
            m = cal["method"]
            if m == "platt":
                proba = 1 / (1 + np.exp(-(cal["a"] * logits + cal["b"])))
            elif m == "beta":
                s = np.clip(_sigmoid(logits), 1e-10, 1 - 1e-10)
                logit = cal["a"] * np.log(s) + cal["b"] * np.log(1 - s) + cal["c"]
                proba = np.clip(_sigmoid(logit), 0, 1)
            elif m == "isotonic":
                model_file = Path(cal["model_file"]).name  # sanitize path
                bst = lgb.Booster(model_file=str(ARTIFACTS / model_file))
                proba = np.clip(bst.predict(logits.reshape(-1, 1)), 0, 1)
            else:
                raise ValueError(
                    f"Unknown calibration method: {m!r}. "
                    "Supported: platt, beta, isotonic"
                )
        else:
            proba = np.asarray(booster.predict(X), dtype=np.float64)
        return {"pred": (proba > 0.5).astype(np.int64), "proba": proba}

    if task == "multiclass":
        proba = np.asarray(booster.predict(X), dtype=np.float64)
        return {"pred": np.argmax(proba, axis=1).astype(np.int64), "proba": proba}

    raise ValueError(f"Unknown task: {task}")


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Equivalence check
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _compare(
    name: str, actual: np.ndarray, expected: np.ndarray, rtol: float
) -> bool:
    """Compare two arrays and log result."""
    if len(actual) != len(expected):
        log.error(
            "%s: FAIL — row count mismatch (codegen=%d, ref=%d)",
            name, len(actual), len(expected),
        )
        return False
    try:
        np.testing.assert_allclose(actual, expected, rtol=rtol)
        log.info("%s: PASS (rtol=%s, n=%d)", name, rtol, len(actual))
        return True
    except AssertionError as e:
        log.error("%s: FAIL\\n%s", name, e)
        return False


def check_equivalence(
    df: pd.DataFrame,
    ref: pd.DataFrame,
    *,
    rtol: float = 1e-7,
) -> bool:
    """Compare codegen predictions against reference."""
    if "pred" not in ref.columns:
        raise ValueError(
            f\'"pred" column not found in reference file. \'
            f"Available columns: {list(ref.columns)}"
        )

    result = _predict(df)
    ok = True

    # Check pred
    if not _compare("pred", result["pred"], ref["pred"].values, rtol):
        ok = False

    # Check proba
    proba = result["proba"]
    if proba is not None:
        if proba.ndim == 1 and "proba" in ref.columns:
            # Binary: single proba column
            if not _compare("proba", proba, ref["proba"].values, rtol):
                ok = False
        elif proba.ndim == 2:
            # Multiclass: proba_0, proba_1, ...
            for i in range(proba.shape[1]):
                col = f"proba_{i}"
                if col in ref.columns:
                    if not _compare(col, proba[:, i], ref[col].values, rtol):
                        ok = False

    return ok


def main() -> None:
    p = argparse.ArgumentParser(
        description="Verify codegen predictions match LizyML reference"
    )
    p.add_argument("data", help="CSV or Parquet input data")
    p.add_argument("--reference", "-r", required=True,
                   help="Reference predictions CSV (from LizyML)")
    p.add_argument("--rtol", type=float, default=1e-7,
                   help="Relative tolerance (default: 1e-7)")
    args = p.parse_args()

    data_path = Path(args.data)
    df = (pd.read_parquet(data_path) if data_path.suffix == ".parquet"
          else pd.read_csv(data_path))
    ref = pd.read_csv(args.reference)

    ok = check_equivalence(df, ref, rtol=args.rtol)
    if ok:
        log.info("All checks PASSED.")
    else:
        log.error("Some checks FAILED.")
        sys.exit(1)


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

_REQUIREMENTS_TXT = """\
lightgbm>=4.0
numpy
pandas
scikit-learn
scipy
"""


def render_train_py() -> str:
    """Return the train.py source code."""
    return _TRAIN_PY


def render_predict_py() -> str:
    """Return the predict.py source code."""
    return _PREDICT_PY


def render_requirements_txt() -> str:
    """Return the requirements.txt content."""
    return _REQUIREMENTS_TXT


def render_test_equivalence_py() -> str:
    """Return the test_equivalence.py source code."""
    return _TEST_EQUIVALENCE_PY
