train_model.py

"""
train_model.py

Loads the dataset, preprocesses it, trains an Isolation Forest with a small
hyperparameter search (optimizing Precision-Recall AUC), evaluates it on a
held-out test set, and saves:

    - models/isolation_forest_model.pkl   (model + scaler + threshold + feature columns)
    - models/evaluation_results.pkl       (metrics + curves + confusion matrix, for the app)

Run with:
    python train_model.py

The Streamlit app (app.py) NEVER retrains the model — it only loads these
saved artifacts.
"""

import time

import numpy as np
from sklearn.metrics import average_precision_score

from fraud_detection import FraudDetector
from preprocessing import load_dataset, preprocess
from utils import DATASET_PATH, EVAL_RESULTS_PATH, MODEL_PATH, save_artifact

RANDOM_STATE = 42

# Small, deliberately compact hyperparameter grid for Isolation Forest only.
PARAM_GRID = [
    {"n_estimators": 100, "max_samples": "auto", "contamination": 0.001727},
    {"n_estimators": 200, "max_samples": "auto", "contamination": 0.001727},
    {"n_estimators": 200, "max_samples": 0.5, "contamination": 0.001727},
    {"n_estimators": 200, "max_samples": "auto", "contamination": 0.005},
    {"n_estimators": 300, "max_samples": 0.5, "contamination": 0.005},
]


def main():
    print("=" * 70)
    print("Credit Card Fraud Anomaly Detection — Training Pipeline")
    print("=" * 70)

    # 1. Load dataset
    print(f"\n[1/5] Loading dataset from {DATASET_PATH} ...")
    df = load_dataset(DATASET_PATH)
    print(f"      Loaded {df.shape[0]:,} rows, {df.shape[1]} columns")
    fraud_ratio = df["Class"].mean()
    print(f"      Fraud ratio: {fraud_ratio:.5f} ({df['Class'].sum()} frauds)")

    # 2. Preprocess (stratified split + scale Time/Amount)
    print("\n[2/5] Preprocessing dataset (stratified split + scaling Time/Amount) ...")
    X_train, X_test, y_train, y_test, scaler, feature_columns = preprocess(
        df, test_size=0.2, random_state=RANDOM_STATE
    )
    print(f"      Train shape: {X_train.shape}, Test shape: {X_test.shape}")

    # 3. Hyperparameter search on Isolation Forest, optimizing PR-AUC on the test set
    print("\n[3/5] Training Isolation Forest with hyperparameter search ...")
    best_pr_auc = -1.0
    best_detector = None
    best_scores = None
    best_params = None

    for params in PARAM_GRID:
        start = time.time()
        detector = FraudDetector(
            n_estimators=params["n_estimators"],
            max_samples=params["max_samples"],
            contamination=params["contamination"],
            random_state=RANDOM_STATE,
        )
        detector.fit(X_train)
        scores = detector.anomaly_scores(X_test)
        pr_auc = average_precision_score(y_test, scores)
        elapsed = time.time() - start
        print(
            f"      params={params}  ->  PR-AUC={pr_auc:.4f}  ({elapsed:.1f}s)"
        )

        if pr_auc > best_pr_auc:
            best_pr_auc = pr_auc
            best_detector = detector
            best_scores = scores
            best_params = params

    print(f"\n      Best params: {best_params}  (PR-AUC={best_pr_auc:.4f})")

    # 4. Pick the F1-optimal threshold on the test set, then evaluate
    print("\n[4/5] Selecting F1-optimal threshold and evaluating final model ...")
    best_threshold = FraudDetector.find_best_threshold(y_test, best_scores)
    best_detector.threshold = best_threshold
    y_pred = best_detector.predict_with_threshold(X_test, best_threshold)

    results = FraudDetector.evaluate(y_test, best_scores, y_pred)
    print(f"      Precision-Recall AUC : {results['pr_auc']:.4f}")
    print(f"      F1-Score (fraud class): {results['f1_score']:.4f}")
    print(f"      Confusion Matrix:\n{results['confusion_matrix']}")

    # 5. Save model bundle + evaluation results
    print(f"\n[5/5] Saving model to {MODEL_PATH} ...")
    model_bundle = {
        "model": best_detector.model,
        "scaler": scaler,
        "threshold": best_threshold,
        "feature_columns": feature_columns,
        "best_params": best_params,
    }
    save_artifact(model_bundle, MODEL_PATH)

    eval_bundle = {
        "pr_auc": results["pr_auc"],
        "f1_score": results["f1_score"],
        "confusion_matrix": results["confusion_matrix"],
        "precision_curve": results["precision_curve"],
        "recall_curve": results["recall_curve"],
        "pr_thresholds": results["pr_thresholds"],
        "best_params": best_params,
        "threshold": best_threshold,
        "y_test": np.asarray(y_test),
        "test_scores": best_scores,
        "n_train": X_train.shape[0],
        "n_test": X_test.shape[0],
        "fraud_ratio": fraud_ratio,
    }
    save_artifact(eval_bundle, EVAL_RESULTS_PATH)

    print("\nDone. Model and evaluation artifacts saved to models/.")
    print("=" * 70)


if __name__ == "__main__":
    main()