churn/models.py

"""Model pipeline factory, candidate registry, and cross-validated leaderboard."""

from __future__ import annotations

import json
import time
from collections import OrderedDict
from pathlib import Path
from typing import Optional

import mlflow
import numpy as np
import pandas as pd
from catboost import CatBoostClassifier
from imblearn.over_sampling import SMOTENC
from imblearn.pipeline import Pipeline as ImbPipeline
from lightgbm import LGBMClassifier
from sklearn.dummy import DummyClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_validate
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import Pipeline
from xgboost import XGBClassifier

from churn.config import settings
from churn.data import get_splits
from churn.features import CT_CATEGORICAL, build_preprocessor

SEED: int = settings.random_seed

# Metrics reported in the leaderboard.
# neg_brier_score is negated by sklearn; we flip the sign when assembling results.
_SCORERS: dict[str, str] = {
    "pr_auc": "average_precision",
    "roc_auc": "roc_auc",
    "brier": "neg_brier_score",
}

# ---------------------------------------------------------------------------
# Pipeline factory
# ---------------------------------------------------------------------------


def build_model_pipeline(estimator) -> Pipeline:
    """Return a full end-to-end Pipeline: preprocessor then *estimator*.

    Reused by every later step (tuning, training, serving) so the same
    preprocessing is always applied in the same order.
    """
    return Pipeline([
        ("preprocessor", build_preprocessor()),
        ("model", estimator),
    ])


# ---------------------------------------------------------------------------
# Candidate model registry
# ---------------------------------------------------------------------------


def get_candidate_models() -> OrderedDict:
    """Return an ordered dict of {name: unfitted estimator}.

    Rules: library defaults + fixed seed + quiet logging only.
    No class_weight / scale_pos_weight / resampling — this is the pure
    out-of-the-box baseline comparison. Imbalance handling comes later.

    MLPClassifier is the optional neural baseline. On tabular data it is
    expected to underperform GBMs; that result is itself informative.
    LogisticRegression gets max_iter=2000 to ensure convergence on 54 features.
    """
    return OrderedDict([
        ("dummy",    DummyClassifier(strategy="prior")),
        ("logreg",   LogisticRegression(max_iter=2000, random_state=SEED)),
        ("xgboost",  XGBClassifier(eval_metric="logloss", random_state=SEED, n_jobs=-1)),
        ("lightgbm", LGBMClassifier(random_state=SEED, n_jobs=-1, verbose=-1)),
        ("catboost", CatBoostClassifier(random_seed=SEED, verbose=0)),
        ("mlp",      MLPClassifier(
            hidden_layer_sizes=(64, 32),
            max_iter=400,
            early_stopping=True,
            random_state=SEED,
        )),
    ])


# ---------------------------------------------------------------------------
# Cross-validated leaderboard
# ---------------------------------------------------------------------------


def run_leaderboard(
    cv: int = 5,
    sample_frac: Optional[float] = None,
    log_to_mlflow: bool = True,
    tracking_uri: Optional[str] = None,
    experiment_name: str = "churn-leaderboard",
    models: Optional[dict] = None,
    reports_dir: Optional[Path] = None,
) -> pd.DataFrame:
    """Cross-validate every candidate on the TRAIN split; return a sorted leaderboard.

    Parameters
    ----------
    cv:
        Number of StratifiedKFold folds.
    sample_frac:
        If set, take a stratified subsample of X_train for fast testing.
    log_to_mlflow:
        Whether to log runs to MLflow. Set False in tests to keep them fast.
    tracking_uri:
        Override the MLflow tracking URI (defaults to settings.mlflow_tracking_uri).
        Pass a tmp_path string in tests to avoid touching the real store.
    experiment_name:
        MLflow experiment name for the parent run.
    models:
        Override the candidate dict (subset for targeted tests).

    Returns
    -------
    DataFrame sorted by pr_auc_mean descending, one row per model.
    Also writes reports/leaderboard.csv.
    """
    X_train, _, y_train, _ = get_splits()

    # Optional stratified subsample for fast CI / unit testing
    if sample_frac is not None:
        from sklearn.model_selection import train_test_split
        X_train, _, y_train, _ = train_test_split(
            X_train, y_train,
            train_size=sample_frac,
            stratify=y_train,
            random_state=SEED,
        )

    candidates: dict = models if models is not None else get_candidate_models()
    cv_splitter = StratifiedKFold(n_splits=cv, shuffle=True, random_state=SEED)

    rows: list[dict] = []

    # --- MLflow parent run setup ---
    if log_to_mlflow:
        uri = tracking_uri or settings.mlflow_tracking_uri
        mlflow.set_tracking_uri(uri)
        mlflow.set_experiment(experiment_name)
        parent_run = mlflow.start_run(run_name="leaderboard")
        parent_run_id = parent_run.info.run_id
    else:
        parent_run_id = None

    try:
        for name, estimator in candidates.items():
            pipe = build_model_pipeline(estimator)

            t0 = time.perf_counter()
            cv_results = cross_validate(
                pipe,
                X_train,
                y_train,
                cv=cv_splitter,
                scoring={k: v for k, v in zip(_SCORERS.keys(), _SCORERS.values())},
                return_train_score=False,
                n_jobs=1,  # avoid nested parallelism conflicts
            )
            elapsed = time.perf_counter() - t0

            pr_auc_scores  = cv_results["test_pr_auc"]
            roc_auc_scores = cv_results["test_roc_auc"]
            # neg_brier_score → flip sign so 0 is perfect, 1 is worst
            brier_scores = -cv_results["test_brier"]

            row = {
                "model":          name,
                "pr_auc_mean":    float(np.mean(pr_auc_scores)),
                "pr_auc_std":     float(np.std(pr_auc_scores)),
                "roc_auc_mean":   float(np.mean(roc_auc_scores)),
                "roc_auc_std":    float(np.std(roc_auc_scores)),
                "brier_mean":     float(np.mean(brier_scores)),
                "brier_std":      float(np.std(brier_scores)),
                "fit_time_mean":  float(elapsed / cv),
            }
            rows.append(row)

            if log_to_mlflow:
                with mlflow.start_run(
                    run_name=name,
                    nested=True,
                    parent_run_id=parent_run_id,
                ):
                    mlflow.log_param("model_name", name)
                    mlflow.log_param("cv_folds", cv)
                    mlflow.log_param("sample_frac", sample_frac)
                    mlflow.log_metric("pr_auc_mean",  row["pr_auc_mean"])
                    mlflow.log_metric("pr_auc_std",   row["pr_auc_std"])
                    mlflow.log_metric("roc_auc_mean", row["roc_auc_mean"])
                    mlflow.log_metric("roc_auc_std",  row["roc_auc_std"])
                    mlflow.log_metric("brier_mean",   row["brier_mean"])
                    mlflow.log_metric("brier_std",    row["brier_std"])
                    mlflow.log_metric("fit_time_mean_s", row["fit_time_mean"])

    finally:
        if log_to_mlflow:
            mlflow.end_run()  # close parent run

    leaderboard = (
        pd.DataFrame(rows)
        .sort_values("pr_auc_mean", ascending=False)
        .reset_index(drop=True)
    )

    # Persist leaderboard artifact
    _reports = Path(reports_dir) if reports_dir is not None else Path("reports")
    _reports.mkdir(exist_ok=True)
    leaderboard_path = _reports / "leaderboard.csv"
    leaderboard.to_csv(leaderboard_path, index=False)

    if log_to_mlflow:
        mlflow.set_tracking_uri(uri)
        with mlflow.start_run(run_id=parent_run_id):
            mlflow.log_artifact(str(leaderboard_path), artifact_path="leaderboard")
        # log leaderboard metadata as JSON on parent for easy API access
        with mlflow.start_run(run_id=parent_run_id):
            mlflow.log_dict(
                json.loads(leaderboard.to_json(orient="records")),
                "leaderboard/summary.json",
            )

    return leaderboard


# ---------------------------------------------------------------------------
# Imbalance experiment
# ---------------------------------------------------------------------------

# Ordered from simplest to most complex — used for tiebreaking in recommendation.
_IMBALANCE_STRATEGIES: list[str] = ["none", "scale_pos_weight", "smotenc"]

# Note: "smotenc_tomek" (SMOTETomek with SMOTENC) is intentionally omitted.
# SMOTETomek's TomekLinks component receives the SMOTENC output, which is a
# mixed-type DataFrame (string categoricals). TomekLinks uses NearestNeighbors
# and requires a numeric array; it fails on string columns. To keep the pipeline
# correct and readable, the Tomek undersampling is excluded here.


def build_imbalance_pipeline(strategy: str, spw: float = 1.0) -> Pipeline | ImbPipeline:
    """Return a full end-to-end pipeline for the given imbalance strategy.

    All strategies use XGBClassifier with the SAME hyperparameters as the
    leaderboard baseline (eval_metric='logloss', random_state=SEED, n_jobs=-1).
    The only variable is how class imbalance is addressed.

    Parameters
    ----------
    strategy : {"none", "scale_pos_weight", "smotenc"}
        none
            Plain preprocessor + XGBClassifier. Reference baseline.
        scale_pos_weight
            Plain preprocessor + XGBClassifier(..., scale_pos_weight=spw).
            Stratified K-fold preserves the class ratio in every fold, so a
            single train-level spw is representative across folds.
        smotenc
            imblearn Pipeline: [fe → SMOTENC → ct → XGBClassifier].
            SMOTENC is categorical-aware: synthetic minority samples for the
            categorical columns are drawn from the observed category set, not
            interpolated numerically. The sampler fires only during fit (on
            each fold's training portion); the validation fold is always
            un-resampled, so CV scores reflect real held-out performance.
    spw : float
        Positive-class weight = #negatives / #positives on y_train (≈2.77).
        Ignored by "none" and "smotenc".
    """
    def _xgb() -> XGBClassifier:
        return XGBClassifier(eval_metric="logloss", random_state=SEED, n_jobs=-1)

    if strategy == "none":
        return build_model_pipeline(_xgb())

    if strategy == "scale_pos_weight":
        return build_model_pipeline(
            XGBClassifier(
                eval_metric="logloss",
                random_state=SEED,
                n_jobs=-1,
                scale_pos_weight=spw,
            )
        )

    if strategy == "smotenc":
        # Pull the two inner steps out of a fresh preprocessor so each
        # pipeline call gets independent, unfitted transformer objects.
        prep = build_preprocessor()
        fe_step = prep.named_steps["fe"]
        ct_step = prep.named_steps["ct"]
        # CT_CATEGORICAL lists the 17 columns that are categorical in the
        # post-FE frame (16 originals + tenure_bucket). SMOTENC draws
        # synthetic values for those columns from observed category sets;
        # the 6 numeric columns are interpolated normally.
        smotenc = SMOTENC(categorical_features=CT_CATEGORICAL, random_state=SEED)
        return ImbPipeline([
            ("fe", fe_step),
            ("smotenc", smotenc),
            ("ct", ct_step),
            ("model", _xgb()),
        ])

    raise ValueError(
        f"Unknown strategy {strategy!r}. Choose from: {_IMBALANCE_STRATEGIES}"
    )


def run_imbalance_experiment(
    cv: int = 5,
    sample_frac: Optional[float] = None,
    log_to_mlflow: bool = True,
    tracking_uri: Optional[str] = None,
    experiment_name: str = "churn-imbalance",
    strategies: Optional[list[str]] = None,
    reports_dir: Optional[Path] = None,
) -> pd.DataFrame:
    """Cross-validate XGBoost under different imbalance-handling strategies.

    Resampling occurs INSIDE each CV fold (imblearn Pipeline applies the
    sampler only on each fold's training portion). Validation-fold metrics
    reflect un-resampled data throughout.

    Parameters
    ----------
    cv : int
        Number of StratifiedKFold folds.
    sample_frac : float | None
        Stratified subsample fraction for fast testing (None = full train set).
    log_to_mlflow : bool
        Whether to log runs to MLflow.
    tracking_uri : str | None
        Override tracking URI (defaults to settings.mlflow_tracking_uri).
    experiment_name : str
        MLflow experiment name.
    strategies : list[str] | None
        Subset of _IMBALANCE_STRATEGIES to run (None = all).

    Returns
    -------
    DataFrame sorted by pr_auc_mean descending, one row per strategy.
    Also writes reports/imbalance_experiment.csv.
    """
    X_train, _, y_train, _ = get_splits()

    if sample_frac is not None:
        from sklearn.model_selection import train_test_split as _tts

        X_train, _, y_train, _ = _tts(
            X_train, y_train,
            train_size=sample_frac,
            stratify=y_train,
            random_state=SEED,
        )

    n_neg = int((y_train == 0).sum())
    n_pos = int((y_train == 1).sum())
    spw = float(n_neg / n_pos)

    run_strategies = strategies if strategies is not None else _IMBALANCE_STRATEGIES
    cv_splitter = StratifiedKFold(n_splits=cv, shuffle=True, random_state=SEED)

    rows: list[dict] = []

    if log_to_mlflow:
        uri = tracking_uri or settings.mlflow_tracking_uri
        mlflow.set_tracking_uri(uri)
        mlflow.set_experiment(experiment_name)
        parent_run = mlflow.start_run(run_name="imbalance-experiment")
        parent_run_id = parent_run.info.run_id
    else:
        parent_run_id = None

    try:
        for strategy in run_strategies:
            pipe = build_imbalance_pipeline(strategy, spw=spw)

            t0 = time.perf_counter()
            cv_results = cross_validate(
                pipe,
                X_train,
                y_train,
                cv=cv_splitter,
                scoring={k: v for k, v in _SCORERS.items()},
                return_train_score=False,
                n_jobs=1,
            )
            elapsed = time.perf_counter() - t0

            row = {
                "strategy": strategy,
                "pr_auc_mean": float(np.mean(cv_results["test_pr_auc"])),
                "pr_auc_std": float(np.std(cv_results["test_pr_auc"])),
                "roc_auc_mean": float(np.mean(cv_results["test_roc_auc"])),
                "roc_auc_std": float(np.std(cv_results["test_roc_auc"])),
                "brier_mean": float(np.mean(-cv_results["test_brier"])),
                "brier_std": float(np.std(-cv_results["test_brier"])),
                "fit_time_mean": float(elapsed / cv),
            }
            rows.append(row)

            if log_to_mlflow:
                with mlflow.start_run(
                    run_name=strategy,
                    nested=True,
                    parent_run_id=parent_run_id,
                ):
                    mlflow.log_param("strategy", strategy)
                    mlflow.log_param("spw", round(spw, 4))
                    mlflow.log_param("cv_folds", cv)
                    mlflow.log_param("sample_frac", sample_frac)
                    mlflow.log_metric("pr_auc_mean", row["pr_auc_mean"])
                    mlflow.log_metric("pr_auc_std", row["pr_auc_std"])
                    mlflow.log_metric("roc_auc_mean", row["roc_auc_mean"])
                    mlflow.log_metric("roc_auc_std", row["roc_auc_std"])
                    mlflow.log_metric("brier_mean", row["brier_mean"])
                    mlflow.log_metric("brier_std", row["brier_std"])

    finally:
        if log_to_mlflow:
            mlflow.end_run()

    result = (
        pd.DataFrame(rows)
        .sort_values("pr_auc_mean", ascending=False)
        .reset_index(drop=True)
    )

    _reports = Path(reports_dir) if reports_dir is not None else Path("reports")
    _reports.mkdir(exist_ok=True)
    result_path = _reports / "imbalance_experiment.csv"
    result.to_csv(result_path, index=False)

    # --- Recommendation ---
    # Rule: pick best PR-AUC. If any simpler strategy is within 1 std of the
    # best, prefer it if its Brier is no worse (simpler = less preprocessing,
    # less variance at serving time).
    best_pr = result["pr_auc_mean"].max()
    best_pr_std = result.loc[result["pr_auc_mean"].idxmax(), "pr_auc_std"]
    simplicity = {s: i for i, s in enumerate(_IMBALANCE_STRATEGIES)}

    candidates = result[result["pr_auc_mean"] >= best_pr - best_pr_std].copy()
    candidates["_rank"] = candidates["strategy"].map(
        lambda s: (candidates["brier_mean"].min() - candidates.loc[
            candidates["strategy"] == s, "brier_mean"
        ].values[0], -simplicity.get(s, 99))
    )
    # Sort: lower Brier first, then simpler strategy
    candidates = candidates.sort_values(
        ["brier_mean", "strategy"],
        key=lambda col: col if col.name == "brier_mean" else col.map(simplicity),
    )
    recommended = candidates.iloc[0]["strategy"]
    rec_row = candidates.iloc[0]
    best_row = result.loc[result["pr_auc_mean"].idxmax()]

    if recommended == best_row["strategy"]:
        rationale = (
            f"'{recommended}' has the best PR-AUC "
            f"({rec_row['pr_auc_mean']:.4f} ± {rec_row['pr_auc_std']:.4f})."
        )
    else:
        pr_gap = best_row["pr_auc_mean"] - rec_row["pr_auc_mean"]
        rationale = (
            f"PR-AUC gap vs best ({pr_gap:.4f}) is within 1 std "
            f"({best_pr_std:.4f}); '{recommended}' chosen for better/equal "
            f"calibration (Brier {rec_row['brier_mean']:.4f} vs "
            f"{best_row['brier_mean']:.4f}) and lower complexity."
        )

    print("\n=== Imbalance Experiment Results ===")
    print(result.to_string(index=False))
    print(f"\nRecommended strategy : {recommended}")
    print(f"Rationale            : {rationale}")
    print(f"scale_pos_weight used: {spw:.4f}  ({n_neg} neg / {n_pos} pos)")

    if log_to_mlflow:
        uri = tracking_uri or settings.mlflow_tracking_uri
        mlflow.set_tracking_uri(uri)
        with mlflow.start_run(run_id=parent_run_id):
            mlflow.log_artifact(str(result_path), artifact_path="imbalance")
            mlflow.log_dict(
                json.loads(result.to_json(orient="records")),
                "imbalance/summary.json",
            )
            mlflow.log_param("recommended_strategy", recommended)

    return result