training/train.py

"""Feature-engineered model comparison and artifact training (MODEL-01..06).

Loads PaySim, builds features via the shared `app.features` module,
splits by `step` (time-respecting -- MODEL-03), compares class-imbalance
handling strategies (FEAT-04) and three model families (MODEL-01), and
writes a results report for user review (MODEL-05) before any model is
promoted to the versioned artifact (MODEL-06).

This script is intentionally split into two stages:

    python -m training.train compare
        Runs the imbalance-strategy comparison and the model comparison,
        writes reports/MODEL_COMPARISON.md + reports/model_charts/*.png,
        and prints a summary. Does NOT save a model artifact -- MODEL-05
        requires results to be shown to the user before a final model is
        selected.

    python -m training.train finalize --model <name> --threshold <t>
        Retrains the chosen model on the full train split and saves
        `model_v1.pkl` + `model_v1.meta.json` (MODEL-06), using the
        user-selected model name and alert threshold.
"""

from __future__ import annotations

import argparse
import gc
import glob
import json
import platform
import sys
from datetime import datetime, timezone
from pathlib import Path

import joblib
import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import sklearn
import xgboost
from imblearn.over_sampling import SMOTE
from imblearn.pipeline import Pipeline as ImbPipeline
from imblearn.under_sampling import RandomUnderSampler
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    average_precision_score,
    f1_score,
    precision_recall_curve,
    precision_score,
    recall_score,
    roc_auc_score,
)
from xgboost import XGBClassifier

from app.features import (
    FEATURE_COLUMNS,
    FEATURE_COLUMNS_CONSERVATIVE,
    engineer_features,
)

REPORTS_DIR = Path("reports")
CHARTS_DIR = REPORTS_DIR / "model_charts"
MODEL_DIR = Path("models")

CANDIDATE_THRESHOLDS = [0.3, 0.5, 0.7]
TRAIN_FRACTION = 0.8
RANDOM_STATE = 42

VARIANTS = {
    "full": FEATURE_COLUMNS,
    "conservative": FEATURE_COLUMNS_CONSERVATIVE,
}


def _report_path(variant: str) -> Path:
    return REPORTS_DIR / f"MODEL_COMPARISON_{variant}.md"


def _results_json_path(variant: str) -> Path:
    return REPORTS_DIR / f"model_results_{variant}.json"


def _pr_chart_name(variant: str) -> str:
    return f"pr_curves_{variant}.png"


def _find_dataset() -> Path:
    candidates = sorted(glob.glob("data/*.csv"))
    if not candidates:
        raise FileNotFoundError("No CSV found under data/.")
    return Path(candidates[0])


def load_features() -> pd.DataFrame:
    df = pd.read_csv(_find_dataset())
    return engineer_features(df)


def time_respecting_split(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    """MODEL-03: split by `step`, never shuffled.

    All rows with step <= the cutoff step go to train, all later rows go
    to test. Because `step` is simulated hours and fraud/legit transactions
    for the same account are already ordered by step, this also prevents
    an account's later (test) transaction from being used to help predict
    its own earlier (train) transaction -- the only direction entity
    leakage could occur here, since features never look forward in step
    order (see app/features.py `_velocity_features`).
    """
    cutoff_step = df["step"].quantile(TRAIN_FRACTION, interpolation="lower")
    train = df[df["step"] <= cutoff_step]
    test = df[df["step"] > cutoff_step]
    return train, test


def _xy(df: pd.DataFrame, feature_columns: list[str]) -> tuple[pd.DataFrame, pd.Series]:
    # float32 halves the in-memory footprint of the feature matrix versus
    # pandas' default float64, which matters at PaySim's ~6.4M-row scale on
    # a memory-constrained host.
    X = df[feature_columns].astype("float32")
    y = df["isFraud"].astype("int8")
    return X, y


# Capped so a single RandomForest fit stays bounded on a memory-constrained
# host: unlimited-depth trees over 5M+ rows can each consume gigabytes.
# max_depth/min_samples_leaf trade a small amount of accuracy for a fit that
# reliably completes; n_jobs is fixed (not -1) so parallel fits don't each
# try to claim every core's worth of memory at once.
RF_N_ESTIMATORS = 100
RF_MAX_DEPTH = 16
RF_MIN_SAMPLES_LEAF = 5
RF_N_JOBS = 4

# SMOTE's default sampling_strategy="auto" oversamples the minority class up
# to full 1:1 parity with the majority -- at PaySim's ~0.13% fraud rate that
# would balloon a ~5.1M-row training fold to ~10M rows. 0.1 (minority raised
# to 10% of majority count) is a common practical middle ground that keeps
# the resampled set memory-tractable while still meaningfully rebalancing.
SMOTE_SAMPLING_STRATEGY = 0.1


def _make_rf(**overrides) -> RandomForestClassifier:
    params = dict(
        n_estimators=RF_N_ESTIMATORS,
        max_depth=RF_MAX_DEPTH,
        min_samples_leaf=RF_MIN_SAMPLES_LEAF,
        n_jobs=RF_N_JOBS,
        random_state=RANDOM_STATE,
    )
    params.update(overrides)
    return RandomForestClassifier(**params)


def compare_imbalance_strategies(
    X_train: pd.DataFrame, y_train: pd.Series, X_test: pd.DataFrame, y_test: pd.Series
) -> dict:
    """FEAT-04: SMOTE vs class-weight vs undersampling, resampling only the
    training fold, via `imblearn.Pipeline` so `fit_resample` never touches
    the held-out test fold.
    """
    strategies = {
        "class_weight": ImbPipeline(
            [("clf", _make_rf(class_weight="balanced"))]
        ),
        "smote": ImbPipeline(
            [
                (
                    "resample",
                    SMOTE(
                        sampling_strategy=SMOTE_SAMPLING_STRATEGY,
                        random_state=RANDOM_STATE,
                    ),
                ),
                ("clf", _make_rf()),
            ]
        ),
        "undersample": ImbPipeline(
            [
                ("resample", RandomUnderSampler(random_state=RANDOM_STATE)),
                ("clf", _make_rf()),
            ]
        ),
    }

    results = {}
    for name, pipeline in strategies.items():
        pipeline.fit(X_train, y_train)
        proba = pipeline.predict_proba(X_test)[:, 1]
        results[name] = {
            "pr_auc": float(average_precision_score(y_test, proba)),
            "roc_auc": float(roc_auc_score(y_test, proba)),
        }
        del pipeline
        gc.collect()
    return results


def train_and_evaluate_models(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    X_test: pd.DataFrame,
    y_test: pd.Series,
    resample_strategy: str,
    model_names: list[str] | None = None,
    keep_fitted: bool = False,
) -> dict:
    """MODEL-01 / MODEL-02: train LogReg/RF/XGB, report PR-AUC/ROC-AUC and
    precision/recall/F1 at each candidate threshold (MODEL-04).

    `model_names` restricts which of the three model families are actually
    built -- `run_finalize` passes a single name so it doesn't pay the
    memory/time cost of also fitting the two models the user didn't pick.
    `keep_fitted` controls whether fitted pipelines are retained in the
    returned dict (`run_compare` never needs them, so leaving this False
    there lets each pipeline be freed right after scoring instead of all
    three living in memory simultaneously).
    """

    def _resampler():
        if resample_strategy == "smote":
            return SMOTE(
                sampling_strategy=SMOTE_SAMPLING_STRATEGY, random_state=RANDOM_STATE
            )
        if resample_strategy == "undersample":
            return RandomUnderSampler(random_state=RANDOM_STATE)
        return None

    models = {
        "logistic_regression": LogisticRegression(
            max_iter=1000,
            class_weight=None if resample_strategy != "class_weight" else "balanced",
            random_state=RANDOM_STATE,
        ),
        "random_forest": _make_rf(
            class_weight=None if resample_strategy != "class_weight" else "balanced",
        ),
        "xgboost": XGBClassifier(
            n_estimators=200,
            max_depth=6,
            learning_rate=0.1,
            eval_metric="aucpr",
            scale_pos_weight=(
                (y_train == 0).sum() / (y_train == 1).sum()
                if resample_strategy == "class_weight"
                else 1
            ),
            random_state=RANDOM_STATE,
            n_jobs=RF_N_JOBS,
        ),
    }

    if model_names is not None:
        models = {name: models[name] for name in model_names}

    results = {}
    fitted = {}
    for name, estimator in models.items():
        resampler = _resampler()
        steps = []
        if resampler is not None:
            steps.append(("resample", resampler))
        steps.append(("clf", estimator))
        pipeline = ImbPipeline(steps)
        pipeline.fit(X_train, y_train)

        proba = pipeline.predict_proba(X_test)[:, 1]
        pr_auc = float(average_precision_score(y_test, proba))
        roc_auc = float(roc_auc_score(y_test, proba))

        threshold_metrics = {}
        for t in CANDIDATE_THRESHOLDS:
            preds = (proba >= t).astype(int)
            threshold_metrics[str(t)] = {
                "precision": float(precision_score(y_test, preds, zero_division=0)),
                "recall": float(recall_score(y_test, preds, zero_division=0)),
                "f1": float(f1_score(y_test, preds, zero_division=0)),
            }

        results[name] = {
            "pr_auc": pr_auc,
            "roc_auc": roc_auc,
            "thresholds": threshold_metrics,
            "proba": proba,  # kept in-memory only, stripped before JSON dump
        }

        if keep_fitted:
            fitted[name] = pipeline
        else:
            del pipeline
            gc.collect()

    return {"results": results, "fitted": fitted}


def _plot_pr_curves(y_test: pd.Series, model_results: dict, variant: str) -> str:
    CHARTS_DIR.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(6, 5))
    for name, r in model_results.items():
        precision, recall, _ = precision_recall_curve(y_test, r["proba"])
        ax.plot(recall, precision, label=f"{name} (PR-AUC={r['pr_auc']:.3f})")
    ax.set_xlabel("Recall")
    ax.set_ylabel("Precision")
    ax.set_title(f"Precision-Recall curves ({variant} feature set)")
    ax.legend()
    path = CHARTS_DIR / _pr_chart_name(variant)
    fig.savefig(path, bbox_inches="tight", dpi=110)
    plt.close(fig)
    return f"model_charts/{path.name}"


def render_comparison_report(
    imbalance_results: dict,
    recommended_strategy: str,
    model_results: dict,
    pr_chart: str,
    train_size: int,
    test_size: int,
    variant: str,
) -> str:
    variant_note = (
        "This run uses the **full** feature set, including "
        "`amount_to_oldbalanceOrg_ratio` and `orig_balance_delta_ratio`. "
        "These two features let a model memorize PaySim's fraud-generation "
        "artifact (amount == oldbalanceOrg in ~98% of fraud rows, 0% of "
        "legit TRANSFER/CASH_OUT rows) rather than learn a generalizable "
        "pattern -- treat PR-AUC here as an upper bound inflated by that "
        "artifact, not a realistic real-world estimate. Compare against "
        "`MODEL_COMPARISON_conservative.md`, which excludes both features."
        if variant == "full"
        else "This run uses the **conservative** feature set: "
        "`amount_to_oldbalanceOrg_ratio` and `orig_balance_delta_ratio` are "
        "excluded because they let a model memorize PaySim's fraud-"
        "generation artifact (amount == oldbalanceOrg in ~98% of fraud "
        "rows, 0% of legit TRANSFER/CASH_OUT rows) rather than learn a "
        "generalizable pattern. These numbers are the more honest estimate "
        "of how the model would perform on data that doesn't share this "
        "synthetic quirk. Compare against `MODEL_COMPARISON_full.md`."
    )

    lines = [
        f"# Phase 2: Feature Engineering & Model Comparison ({variant} feature set)",
        "",
        f"Train rows: {train_size:,} / Test rows: {test_size:,} "
        "(time-respecting split on `step`, MODEL-03).",
        "",
        variant_note,
        "",
        "## Class imbalance strategy comparison (FEAT-04)",
        "",
        "Random Forest baseline, held-out `step`-based test fold, resampling "
        "applied only within the training fold via `imblearn.Pipeline`:",
        "",
        "| Strategy | PR-AUC | ROC-AUC |",
        "|---|---|---|",
    ]
    for name, r in imbalance_results.items():
        marker = " **<- recommended**" if name == recommended_strategy else ""
        lines.append(f"| {name} | {r['pr_auc']:.4f} | {r['roc_auc']:.4f}{marker} |")

    lines += [
        "",
        f"**Recommendation: `{recommended_strategy}`** (highest PR-AUC on the "
        "held-out fold). PR-AUC is used to pick the strategy, not ROC-AUC, "
        "because ROC-AUC is optimistic under extreme class imbalance "
        "(~0.13% fraud) -- PR-AUC is far more sensitive to false positives "
        "against the large negative class.",
        "",
        "## Model comparison (MODEL-01, MODEL-02)",
        "",
        f"All three models below use the `{recommended_strategy}` imbalance "
        "strategy.",
        "",
        "| Model | PR-AUC | ROC-AUC |",
        "|---|---|---|",
    ]
    for name, r in model_results.items():
        lines.append(f"| {name} | {r['pr_auc']:.4f} | {r['roc_auc']:.4f} |")

    lines += [
        "",
        f"![Precision-Recall curves]({pr_chart})",
        "",
        "## Precision / Recall / F1 at candidate thresholds (MODEL-04)",
        "",
    ]
    for name, r in model_results.items():
        lines.append(f"### {name}")
        lines.append("")
        lines.append("| Threshold | Precision | Recall | F1 |")
        lines.append("|---|---|---|---|")
        for t in CANDIDATE_THRESHOLDS:
            m = r["thresholds"][str(t)]
            lines.append(
                f"| {t} | {m['precision']:.4f} | {m['recall']:.4f} | {m['f1']:.4f} |"
            )
        lines.append("")

    lines += [
        "## Next step (MODEL-05 checkpoint)",
        "",
        "These results are presented for review before a final model is "
        "selected. Once you choose a model, a feature-set variant, and an "
        "alert threshold from the tables above (this report and its "
        "counterpart for the other variant), run:",
        "",
        "```",
        "python -m training.train finalize --variant "
        f"{variant} --model <name> --threshold <t>",
        "```",
        "",
        "to retrain the chosen model on the full training split and save "
        "the versioned artifact (`models/model_v1.pkl`).",
        "",
    ]
    return "\n".join(lines)


def run_compare(variant: str) -> None:
    feature_columns = VARIANTS[variant]
    print(f"[{variant}] Loading data and engineering features...", flush=True)
    df = load_features()
    train_df, test_df = time_respecting_split(df)
    train_rows, test_rows = len(train_df), len(test_df)
    X_train, y_train = _xy(train_df, feature_columns)
    X_test, y_test = _xy(test_df, feature_columns)
    # Free the raw/engineered frames now that the feature/label arrays are
    # extracted -- `df`, `train_df`, `test_df` together hold a second,
    # float64, all-columns copy of the ~6.4M-row dataset that nothing below
    # this point needs.
    del df, train_df, test_df
    gc.collect()

    print(f"[{variant}] Train: {train_rows:,} rows / Test: {test_rows:,} rows", flush=True)

    print(f"[{variant}] Comparing class-imbalance strategies...", flush=True)
    imbalance_results = compare_imbalance_strategies(X_train, y_train, X_test, y_test)
    recommended_strategy = max(
        imbalance_results, key=lambda k: imbalance_results[k]["pr_auc"]
    )
    print(f"[{variant}] Recommended imbalance strategy: {recommended_strategy}", flush=True)

    print(f"[{variant}] Training and comparing LogReg / RandomForest / XGBoost...", flush=True)
    outcome = train_and_evaluate_models(
        X_train, y_train, X_test, y_test, recommended_strategy
    )
    model_results = outcome["results"]

    pr_chart = _plot_pr_curves(y_test, model_results, variant)

    REPORTS_DIR.mkdir(parents=True, exist_ok=True)
    report = render_comparison_report(
        imbalance_results,
        recommended_strategy,
        model_results,
        pr_chart,
        train_rows,
        test_rows,
        variant,
    )
    report_path = _report_path(variant)
    report_path.write_text(report)

    # persist metrics (without the raw proba arrays) for the finalize step
    # and for MODEL-06's "metrics logged to a results file" requirement
    serializable = {
        "variant": variant,
        "feature_columns": feature_columns,
        "recommended_imbalance_strategy": recommended_strategy,
        "imbalance_comparison": imbalance_results,
        "model_comparison": {
            name: {k: v for k, v in r.items() if k != "proba"}
            for name, r in model_results.items()
        },
        "train_rows": train_rows,
        "test_rows": test_rows,
        "generated_at": datetime.now(timezone.utc).isoformat(),
    }
    results_path = _results_json_path(variant)
    results_path.write_text(json.dumps(serializable, indent=2))

    print(f"\n[{variant}] Report written to {report_path}")
    print(f"[{variant}] Results JSON written to {results_path}")
    print(f"[{variant}] Chart written to {CHARTS_DIR / _pr_chart_name(variant)}")


def run_finalize(variant: str, model_name: str, threshold: float) -> None:
    results_path = _results_json_path(variant)
    if not results_path.exists():
        raise FileNotFoundError(
            f"{results_path} not found -- run "
            f"`python -m training.train compare --variant {variant}` first."
        )
    prior = json.loads(results_path.read_text())
    strategy = prior["recommended_imbalance_strategy"]
    feature_columns = prior["feature_columns"]
    if model_name not in prior["model_comparison"]:
        raise ValueError(
            f"Unknown model '{model_name}'. Choose one of: "
            f"{list(prior['model_comparison'])}"
        )

    print(
        f"Retraining '{model_name}' ({variant} feature set) with "
        f"'{strategy}' resampling on full data...",
        flush=True,
    )
    df = load_features()
    train_df, test_df = time_respecting_split(df)
    train_rows, test_rows = len(train_df), len(test_df)
    X_train, y_train = _xy(train_df, feature_columns)
    X_test, y_test = _xy(test_df, feature_columns)
    del df, train_df, test_df
    gc.collect()

    outcome = train_and_evaluate_models(
        X_train,
        y_train,
        X_test,
        y_test,
        strategy,
        model_names=[model_name],
        keep_fitted=True,
    )
    pipeline = outcome["fitted"][model_name]
    metrics = {k: v for k, v in outcome["results"][model_name].items() if k != "proba"}
    metrics["selected_threshold"] = threshold
    metrics["threshold_metrics_at_selection"] = metrics["thresholds"][str(threshold)] \
        if str(threshold) in metrics["thresholds"] else None

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    artifact_path = MODEL_DIR / "model_v1.pkl"
    joblib.dump(pipeline, artifact_path)

    meta = {
        "model_name": model_name,
        "feature_variant": variant,
        "imbalance_strategy": strategy,
        "feature_columns": feature_columns,
        "selected_threshold": threshold,
        "metrics": metrics,
        "training_versions": {
            "python": sys.version,
            "platform": platform.platform(),
            "scikit_learn": sklearn.__version__,
            "xgboost": xgboost.__version__,
            "numpy": np.__version__,
            "pandas": pd.__version__,
        },
        "trained_at": datetime.now(timezone.utc).isoformat(),
        "train_rows": train_rows,
        "test_rows": test_rows,
    }
    meta_path = MODEL_DIR / "model_v1.meta.json"
    meta_path.write_text(json.dumps(meta, indent=2, default=str))

    print(f"Saved artifact to {artifact_path}")
    print(f"Saved metadata to {meta_path}")
    print(json.dumps(metrics, indent=2, default=str))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)

    compare_parser = sub.add_parser("compare", help="Run the imbalance + model comparison")
    compare_parser.add_argument(
        "--variant", choices=list(VARIANTS), default="full"
    )

    finalize_parser = sub.add_parser("finalize", help="Save the chosen model")
    finalize_parser.add_argument("--variant", choices=list(VARIANTS), default="full")
    finalize_parser.add_argument("--model", required=True)
    finalize_parser.add_argument("--threshold", required=True, type=float)

    args = parser.parse_args()
    if args.command == "compare":
        run_compare(args.variant)
    elif args.command == "finalize":
        run_finalize(args.variant, args.model, args.threshold)


if __name__ == "__main__":
    main()