src/models/train.py

"""XGBoost churn model: training, evaluation, and artifact persistence.

Run as a module:
    uv run python -m src.models.train
"""

import json
import pickle
from pathlib import Path

import pandas as pd
import xgboost as xgb
from sklearn.metrics import roc_curve

from src.baseline.rule_based import evaluate_baseline
from src.config import (
    ARTIFACTS_DIR,
    AUG_N_WINDOWS,
    AUG_STEP_DAYS,
    CHURN_PERIOD_DAYS,
    DATA_DIR,
    TEST_FRACTION,
    XGBOOST_PARAMS,
)
from src.evaluation.calibration import compute_calibration
from src.evaluation.distribution_shift import compute_adversarial_auc
from src.evaluation.guardrails import run_all_guardrails
from src.evaluation.metrics import compute_metrics
from src.feature_engineering.pipeline import build_feature_matrix
from src.preprocessing.clean import clean_customers, clean_orders
from src.target.churn_target import compute_churn_target
from src.utils.logger import get_logger

logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def get_feature_cols(features: pd.DataFrame) -> list[str]:
    """Return all column names except ``customer_id``."""
    return [c for c in features.columns if c != "customer_id"]


def time_based_split(
    features: pd.DataFrame,
    target: pd.DataFrame,
    customers: pd.DataFrame,
    test_fraction: float = TEST_FRACTION,
) -> tuple[pd.DataFrame, pd.Series, pd.DataFrame, pd.Series, list[str]]:
    """Stratified-by-registration-cohort train / test split – no shuffle.

    Customers are sorted by ``registration_date`` and then every
    ``round(1 / test_fraction)``-th customer (0-indexed) is assigned to the
    test set; the rest form the training set.  This interleaved selection
    ensures both splits cover the same registration-date range, preventing
    the cohort confound that arises from a hard temporal cut (which would
    place all late-joiners in the test set and all early-joiners in train).

    Returns:
        X_train, y_train, X_test, y_test, feature_cols, test_customer_ids
    """
    df = features.merge(target, on="customer_id", how="inner")
    reg_dates = customers[["customer_id", "registration_date"]]
    df = df.merge(reg_dates, on="customer_id", how="left")
    df = df.sort_values("registration_date").reset_index(drop=True)

    every_nth = max(2, round(1.0 / test_fraction))
    test_mask = df.index % every_nth == (every_nth - 1)
    train_df = df[~test_mask]
    test_df = df[test_mask]

    feature_cols = get_feature_cols(features)
    X_train = train_df[feature_cols]
    y_train = train_df["is_churned"]
    X_test = test_df[feature_cols]
    y_test = test_df["is_churned"]
    test_customer_ids = test_df["customer_id"].tolist()

    logger.info(
        f"Split: train={len(X_train):,} (churn={y_train.mean():.2%})  "
        f"test={len(X_test):,} (churn={y_test.mean():.2%})"
    )
    return X_train, y_train, X_test, y_test, feature_cols, test_customer_ids


# ---------------------------------------------------------------------------
# Model training
# ---------------------------------------------------------------------------


def train_xgboost(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    params: dict | None = None,
) -> xgb.XGBClassifier:
    """Train an XGBoost classifier on the provided data.

    Args:
        X_train: Feature matrix.
        y_train: Binary target.
        params: Hyperparameter dict; defaults to ``XGBOOST_PARAMS``.

    Returns:
        Fitted XGBClassifier.
    """
    if params is None:
        params = XGBOOST_PARAMS.copy()

    model = xgb.XGBClassifier(**params)
    model.fit(X_train, y_train, verbose=False)
    logger.info("XGBoost training complete")
    return model


# ---------------------------------------------------------------------------
# Sliding-window data augmentation
# ---------------------------------------------------------------------------


def build_augmented_training_set(
    orders: pd.DataFrame,
    customers: pd.DataFrame,
    train_customer_ids: list,
    feature_end_date: pd.Timestamp,
    feature_cols: list[str],
    n_windows: int = AUG_N_WINDOWS,
    step_days: int = AUG_STEP_DAYS,
    churn_period_days: int = CHURN_PERIOD_DAYS,
) -> tuple[pd.DataFrame, pd.Series]:
    """Generate extra (features, label) rows by sliding the churn window back in time.

    For each window i ∈ {1, …, n_windows}:
      aug_feature_end  = feature_end_date − i × step_days
      aug_t_max        = aug_feature_end  + churn_period_days
      label            = no orders in (aug_feature_end, aug_t_max]
      features         = build_feature_matrix(orders, ..., reference_date=aug_feature_end)

    Only training-set customers are used; test-set customers are never included.
    All features are computed on orders strictly before ``aug_feature_end``, so there
    is no leakage from the churn observation window or from future data.

    Args:
        orders: Full cleaned orders fact table.
        customers: Customer dimension table.
        train_customer_ids: Customer IDs that belong to the training split.
        feature_end_date: The original feature cutoff (= T_max − churn_period_days).
        feature_cols: Column names expected by the model (must match the primary window).
        n_windows: Number of additional windows to generate.
        step_days: Days to step back per window.
        churn_period_days: Length of each churn observation window.

    Returns:
        Tuple of (X_aug, y_aug) — concatenated across all windows.
        Returns empty DataFrames if no eligible customers are found in any window.
    """
    train_id_set = set(train_customer_ids)
    train_customers = customers[customers["customer_id"].isin(train_id_set)]
    X_parts: list[pd.DataFrame] = []
    y_parts: list[pd.Series] = []

    for i in range(1, n_windows + 1):
        aug_feature_end: pd.Timestamp = feature_end_date - pd.Timedelta(days=i * step_days)
        aug_t_max: pd.Timestamp = aug_feature_end + pd.Timedelta(days=churn_period_days)

        # Eligible: placed at least one order on or before the augmented feature cutoff
        pre_window_ids = (
            set(orders[orders["order_date"] <= aug_feature_end]["customer_id"]) & train_id_set
        )
        if not pre_window_ids:
            logger.warning(f"Aug window {i}: no eligible training customers — skipping")
            continue

        # Labels: churned = no orders in the augmented churn window
        active_ids = set(
            orders[(orders["order_date"] > aug_feature_end) & (orders["order_date"] <= aug_t_max)][
                "customer_id"
            ]
        )
        window_customers = train_customers[train_customers["customer_id"].isin(pre_window_ids)]
        target = pd.DataFrame({"customer_id": list(pre_window_ids)})
        target["is_churned"] = (~target["customer_id"].isin(active_ids)).astype(int)

        # Features: strictly before aug_feature_end
        feats = build_feature_matrix(orders, window_customers, aug_feature_end)
        merged = feats.merge(target, on="customer_id", how="inner")
        if merged.empty:
            continue

        X_parts.append(merged[feature_cols])
        y_parts.append(merged["is_churned"].reset_index(drop=True))
        logger.info(
            f"Aug window {i}: feature_end={aug_feature_end.date()}  "
            f"t_max={aug_t_max.date()}  n={len(merged):,}  "
            f"churn={merged['is_churned'].mean():.2%}"
        )

    if not X_parts:
        logger.warning("No augmentation windows produced data; returning empty DataFrames")
        return pd.DataFrame(columns=feature_cols), pd.Series(dtype=int, name="is_churned")

    return (
        pd.concat(X_parts, ignore_index=True),
        pd.concat(y_parts, ignore_index=True),
    )


# ---------------------------------------------------------------------------
# Model persistence
# ---------------------------------------------------------------------------


def save_model(model: xgb.XGBClassifier, path: Path | None = None) -> Path:
    """Pickle the trained model to disk."""
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
    if path is None:
        path = ARTIFACTS_DIR / "xgb_churn_model.pkl"
    with open(path, "wb") as fh:
        pickle.dump(model, fh)
    logger.info(f"Model saved → {path}")
    return path


def load_model(path: Path | None = None) -> xgb.XGBClassifier:
    """Load a previously saved model from disk.

    Raises:
        FileNotFoundError: if the model file does not exist.
    """
    if path is None:
        path = ARTIFACTS_DIR / "xgb_churn_model.pkl"
    if not path.exists():
        raise FileNotFoundError(f"Model not found at {path}. Run `make train` first.")
    with open(path, "rb") as fh:
        return pickle.load(fh)  # nosec B301 - loading trusted local model artifact


_REGULARIZATION_KEYS: frozenset[str] = frozenset({"gamma", "reg_alpha", "min_child_weight"})


def load_best_params() -> dict | None:
    """Load tuned hyperparameters saved by ``tune.py``, if available."""
    params_path = ARTIFACTS_DIR / "best_params.json"
    if params_path.exists():
        with open(params_path) as fh:
            data = json.load(fh)
        logger.info(f"Loaded tuned params (CV ROC-AUC={data['score']:.4f})")
        missing = _REGULARIZATION_KEYS - set(data["params"])
        if missing:
            logger.warning(
                f"Loaded params are missing regularisation keys {missing}. "
                "They were tuned without these constraints and may overfit. "
                f"Delete {params_path} or run `make tune` to re-tune."
            )
        return data["params"]
    logger.info("No tuned params found – using defaults")
    return None


# ---------------------------------------------------------------------------
# Artifact persistence for dashboard
# ---------------------------------------------------------------------------


def save_evaluation_artifacts(
    model: xgb.XGBClassifier,
    metrics: dict,
    calibration: dict,
    guardrails: list[dict],
    X_test: pd.DataFrame,
    y_test: pd.Series,
    feature_cols: list[str],
    baseline_metrics: dict,
    shift_metrics: dict,
) -> None:
    """Write all evaluation artefacts consumed by the dashboard."""
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

    with open(ARTIFACTS_DIR / "eval_metrics.json", "w") as fh:
        json.dump(
            {
                **metrics,
                "baseline_roc_auc": baseline_metrics.get("roc_auc"),
                "baseline_pr_auc": baseline_metrics.get("pr_auc"),
                "baseline_brier_score": baseline_metrics.get("brier_score"),
                "baseline_precision": baseline_metrics.get("precision"),
                "baseline_recall": baseline_metrics.get("recall"),
                "baseline_f1": baseline_metrics.get("f1"),
                "baseline_n_samples": baseline_metrics.get("n_samples"),
            },
            fh,
            indent=2,
        )

    with open(ARTIFACTS_DIR / "calibration.json", "w") as fh:
        json.dump(calibration, fh, indent=2)

    y_prob = model.predict_proba(X_test)[:, 1]
    fpr, tpr, _ = roc_curve(y_test, y_prob)
    with open(ARTIFACTS_DIR / "roc_curve.json", "w") as fh:
        json.dump({"fpr": fpr.tolist(), "tpr": tpr.tolist()}, fh)

    importance = dict(zip(feature_cols, model.feature_importances_.tolist()))
    importance_sorted = dict(sorted(importance.items(), key=lambda kv: kv[1], reverse=True))
    with open(ARTIFACTS_DIR / "feature_importance.json", "w") as fh:
        json.dump(importance_sorted, fh, indent=2)

    with open(ARTIFACTS_DIR / "feature_cols.json", "w") as fh:
        json.dump(feature_cols, fh)

    with open(ARTIFACTS_DIR / "guardrails.json", "w") as fh:
        json.dump(guardrails, fh, indent=2)

    with open(ARTIFACTS_DIR / "distribution_shift.json", "w") as fh:
        json.dump(shift_metrics, fh, indent=2)

    logger.info(f"Evaluation artefacts saved → {ARTIFACTS_DIR}")


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    logger.info("=== Churn Model Training ===")

    orders_raw = pd.read_csv(DATA_DIR / "orders.csv", parse_dates=["order_date", "contract_date"])
    customers_raw = pd.read_csv(
        DATA_DIR / "customers.csv",
        parse_dates=["registration_date", "birth_date", "last_profile_update"],
    )

    orders = clean_orders(orders_raw)
    customers = clean_customers(customers_raw)

    target_df, feature_end_date, t_max = compute_churn_target(orders, customers, CHURN_PERIOD_DAYS)
    features = build_feature_matrix(orders, customers, feature_end_date)

    X_train, y_train, X_test, y_test, feature_cols, test_customer_ids = time_based_split(
        features, target_df, customers
    )

    # --- Distribution shift: can the model tell train from test? ---
    shift_metrics = compute_adversarial_auc(X_train, X_test)

    # --- Sliding-window data augmentation (training customers only, no leakage) ---
    train_customer_ids = [
        cid for cid in target_df["customer_id"].tolist() if cid not in set(test_customer_ids)
    ]
    X_aug, y_aug = build_augmented_training_set(
        orders,
        customers,
        train_customer_ids,
        feature_end_date,
        feature_cols,
        n_windows=AUG_N_WINDOWS,
        step_days=AUG_STEP_DAYS,
        churn_period_days=CHURN_PERIOD_DAYS,
    )
    if not X_aug.empty:
        X_train = pd.concat([X_train, X_aug], ignore_index=True)
        y_train = pd.concat([y_train, y_aug], ignore_index=True)
        logger.info(f"Augmented training set: {len(X_train):,} rows (churn={y_train.mean():.2%})")

    params = load_best_params() or XGBOOST_PARAMS.copy()
    model = train_xgboost(X_train, y_train, params)

    # --- Test-set metrics ---
    y_prob_test = model.predict_proba(X_test)[:, 1]
    metrics = compute_metrics(y_test.to_numpy(), y_prob_test)
    calibration = compute_calibration(y_test.to_numpy(), y_prob_test)

    logger.info(
        f"ML test   ROC-AUC={metrics['roc_auc']:.4f}  "
        f"PR-AUC={metrics['pr_auc']:.4f}  "
        f"Brier={metrics['brier_score']:.4f}"
    )

    # --- Train-set metrics (overfitting detection) ---
    y_prob_train = model.predict_proba(X_train)[:, 1]
    train_metrics = compute_metrics(y_train.to_numpy(), y_prob_train)
    logger.info(
        f"ML train  ROC-AUC={train_metrics['roc_auc']:.4f}  "
        f"gap={train_metrics['roc_auc'] - metrics['roc_auc']:.4f}"
    )

    # --- Baseline evaluated on the **same test set** for a fair comparison ---
    test_features = features[features["customer_id"].isin(test_customer_ids)]
    test_target = target_df[target_df["customer_id"].isin(test_customer_ids)]
    baseline_metrics = evaluate_baseline(test_features, test_target)

    if baseline_metrics["n_samples"] != len(y_test):
        raise RuntimeError(
            f"Sample count mismatch: ML test set has {len(y_test)} rows but "
            f"baseline evaluated on {baseline_metrics['n_samples']}. "
            "Ensure evaluate_baseline receives only test-split customers."
        )

    # --- Detailed side-by-side comparison (same N, same labels) ---
    n = len(y_test)
    cr = metrics["churn_rate"]
    logger.info(f"ML vs Baseline — test set  N={n:,}  churn_rate={cr:.2%}")
    _cmp = [
        ("ROC-AUC", "roc_auc", "higher=better"),
        ("PR-AUC", "pr_auc", "higher=better"),
        ("Brier", "brier_score", "lower=better"),
        ("Precision", "precision", "higher=better"),
        ("Recall", "recall", "higher=better"),
        ("F1", "f1", "higher=better"),
    ]
    for label, key, direction in _cmp:
        ml_v = metrics[key]
        bl_v = baseline_metrics[key]
        delta = ml_v - bl_v
        logger.info(
            f"  {label:<10}  ML={ml_v:.4f}  baseline={bl_v:.4f}  delta={delta:+.4f}  ({direction})"
        )

    predictions_stub = pd.DataFrame({"customer_id": features["customer_id"]})
    guardrails_results = run_all_guardrails(
        metrics,
        calibration,
        predictions_stub,
        customers,
        train_metrics=train_metrics,
        baseline_metrics=baseline_metrics,
        shift_metrics=shift_metrics,
    )
    for g in guardrails_results:
        status = "PASS" if g["passed"] else "FAIL"
        logger.info(f"  [{status}] {g['message']}")

    save_model(model)
    save_evaluation_artifacts(
        model,
        metrics,
        calibration,
        guardrails_results,
        X_test,
        y_test,
        feature_cols,
        baseline_metrics,
        shift_metrics,
    )

    # --- Category-level summary for dashboard (test-set actual vs predicted) ---
    test_df_cats = pd.DataFrame(
        {"customer_id": test_customer_ids, "actual": y_test.to_numpy(), "predicted": y_prob_test}
    )
    test_df_cats = test_df_cats.merge(
        customers[["customer_id", "segment", "marketing_channel", "country"]],
        on="customer_id",
        how="left",
    )
    category_summary: dict = {}
    for raw_col, key in [
        ("segment", "segment"),
        ("marketing_channel", "channel"),
        ("country", "country"),
    ]:
        grp = (
            test_df_cats.groupby(raw_col)
            .agg(
                actual_churn_rate=("actual", "mean"),
                predicted_churn_rate=("predicted", "mean"),
                count=("actual", "count"),
            )
            .reset_index()
            .sort_values(raw_col)
            .rename(columns={raw_col: "category"})
        )
        category_summary[key] = grp.to_dict("records")

    customer_counts: dict = {}
    for raw_col, key in [
        ("segment", "segment"),
        ("marketing_channel", "channel"),
        ("country", "country"),
    ]:
        customer_counts[key] = customers[raw_col].value_counts().to_dict()

    with open(ARTIFACTS_DIR / "category_summary.json", "w") as fh:
        json.dump(category_summary, fh, indent=2)
    with open(ARTIFACTS_DIR / "customer_stats.json", "w") as fh:
        json.dump(customer_counts, fh, indent=2)
    logger.info("Category summary artifacts saved")

    logger.info("=== Training complete ===")