app/training/calibration_analysis.py

"""
Calibration analysis for LightGBM (reviewer priority — Brier alone is not
enough evidence of good calibration).

Produces:
  1. calibration_diagnostics.csv — ECE, calibration slope, calibration
     intercept, Brier score, alongside the reliability-diagram bins.
  2. extended_metrics_table.csv  — Table 3 (all 11 models) extended with
     PR-AUC, Balanced Accuracy, and MCC, computed from the same OOF
     predictions stored in models/training_results.json (no retraining).

Usage:
    python -m app.training.calibration_analysis
"""

from __future__ import annotations

import csv
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from sklearn.calibration import calibration_curve
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    average_precision_score,
    balanced_accuracy_score,
    matthews_corrcoef,
)

MODELS_DIR = Path(__file__).resolve().parents[2] / "models"
TABLES_DIR = Path(__file__).resolve().parents[3] / "docs/academic/paper/real_tables"


def _expected_calibration_error(y_true: np.ndarray, y_prob: np.ndarray, n_bins: int = 10) -> float:
    bin_edges = np.linspace(0.0, 1.0, n_bins + 1)
    bin_indices = np.digitize(y_prob, bin_edges[1:-1])
    ece = 0.0
    n = len(y_true)
    for b in range(n_bins):
        mask = bin_indices == b
        if not np.any(mask):
            continue
        bin_conf = float(np.mean(y_prob[mask]))
        bin_acc = float(np.mean(y_true[mask]))
        ece += (np.sum(mask) / n) * abs(bin_conf - bin_acc)
    return ece


def _calibration_slope_intercept(y_true: np.ndarray, y_prob: np.ndarray) -> tuple[float, float]:
    """Logistic recalibration: y ~ sigmoid(slope * logit(p) + intercept).
    slope=1, intercept=0 is perfect calibration."""
    eps = 1e-6
    p_clipped = np.clip(y_prob, eps, 1 - eps)
    logit_p = np.log(p_clipped / (1 - p_clipped)).reshape(-1, 1)
    lr = LogisticRegression()
    lr.fit(logit_p, y_true)
    slope = float(lr.coef_[0][0])
    intercept = float(lr.intercept_[0])
    return slope, intercept


def run() -> None:
    results_path = MODELS_DIR / "training_results.json"
    with open(results_path, "r", encoding="utf-8") as f:
        training_results = json.load(f)

    # ── Extended metrics table (Table 3 + PR-AUC/BalAcc/MCC) ──
    # training_results.json's per-model entries were saved without the
    # y_true/y_pred/y_prob arrays (train_classifier.py strips those before
    # writing JSON) — so PR-AUC/BalAcc/MCC must come from re-running
    # cross_val_predict, matching evaluate_predictions()'s new metrics.
    # Simpler and fully consistent: reuse train_classifier's train() output
    # in-process is out of scope here; instead recompute from the model
    # pickles + features.csv using the same OOF fold logic as ensemble_model.py.
    print("Computing extended metrics (PR-AUC, Balanced Accuracy, MCC) for all 11 models...")
    from sklearn.model_selection import StratifiedKFold
    from sklearn.preprocessing import StandardScaler
    from sklearn.base import clone
    from sklearn.metrics import roc_auc_score, f1_score, accuracy_score, roc_curve
    import pickle
    import warnings
    from sklearn.exceptions import ConvergenceWarning

    from app.training.evaluate import load_features_csv

    FEATURES_CSV = Path("D:/CrownCode/DataSet/features.csv")
    DL_OOF_NPZ = MODELS_DIR / "dl_oof_probs.npz"
    X, y = load_features_csv(FEATURES_CSV)
    X = np.nan_to_num(X, nan=0.0, posinf=1.0, neginf=-1.0)

    ml_files = {
        "Logistic Regression": "model_logistic_regression.pkl",
        "Random Forest": "model_random_forest.pkl",
        "Gradient Boosting": "model_gradient_boosting.pkl",
        "SVM (RBF)": "model_svm_rbf.pkl",
        "MLP Neural Network": "model_mlp_neural_network.pkl",
        "XGBoost": "model_xgboost.pkl",
        "LightGBM": "model_lightgbm.pkl",
    }
    dl_npz_keys = {
        "Deep MLP (512-256-128-64)": "Deep_MLP_512_256_128_64",
        "1D-CNN": "1D_CNN",
        "Residual MLP (3 blocks)": "Residual_MLP_3_blocks",
        "Attention MLP": "Attention_MLP",
    }

    if not DL_OOF_NPZ.exists():
        raise RuntimeError(f"{DL_OOF_NPZ} not found — run dump_dl_oof.py first.")
    dl_npz = np.load(DL_OOF_NPZ)
    if not np.array_equal(dl_npz["y"], y):
        raise RuntimeError("dl_oof_probs.npz label order does not match features.csv load order.")

    ml_models_raw = {}
    for name, fname in ml_files.items():
        with open(MODELS_DIR / fname, "rb") as f:
            ml_models_raw[name] = pickle.load(f)

    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    fold_assignments = list(cv.split(X, y))
    n = len(y)
    oof_probs: dict[str, np.ndarray] = {name: np.zeros(n) for name in list(ml_files) + list(dl_npz_keys)}
    for name, npz_key in dl_npz_keys.items():
        oof_probs[name] = dl_npz[npz_key]

    for fold_idx, (train_idx, test_idx) in enumerate(fold_assignments, start=1):
        print(f"  Fold {fold_idx}/5 (ML models) ...")
        X_train, y_train = X[train_idx], y[train_idx]
        X_test, y_test = X[test_idx], y[test_idx]

        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train)
        X_test_scaled = scaler.transform(X_test)

        for name in ml_files:
            model = clone(ml_models_raw[name])
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=ConvergenceWarning)
                model.fit(X_train_scaled, y_train)
            oof_probs[name][test_idx] = model.predict_proba(X_test_scaled)[:, 1]

    extended_rows = []
    for name, probs in oof_probs.items():
        fpr, tpr, thr = roc_curve(y, probs)
        threshold = float(thr[np.argmax(tpr - fpr)])
        y_pred = (probs >= threshold).astype(int)
        extended_rows.append({
            "model": name,
            "roc_auc": round(float(roc_auc_score(y, probs)), 4),
            "pr_auc": round(float(average_precision_score(y, probs)), 4),
            "accuracy": round(float(accuracy_score(y, y_pred)), 4),
            "f1": round(float(f1_score(y, y_pred, zero_division=0)), 4),
            "balanced_accuracy": round(float(balanced_accuracy_score(y, y_pred)), 4),
            "mcc": round(float(matthews_corrcoef(y, y_pred)), 4),
            "threshold": round(threshold, 4),
        })

    extended_rows.sort(key=lambda r: -r["roc_auc"])
    TABLES_DIR.mkdir(parents=True, exist_ok=True)
    ext_path = TABLES_DIR / "extended_metrics_table.csv"
    with open(ext_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(extended_rows[0].keys()))
        writer.writeheader()
        writer.writerows(extended_rows)
    print(f"\nExtended metrics table written: {ext_path}")
    for r in extended_rows:
        print(f"  {r['model']:28s} AUC={r['roc_auc']:.4f} PR-AUC={r['pr_auc']:.4f} "
              f"BalAcc={r['balanced_accuracy']:.4f} MCC={r['mcc']:.4f}")

    # ── Calibration diagnostics for LightGBM specifically ──
    lgbm_probs = oof_probs["LightGBM"]
    prob_true, prob_pred = calibration_curve(y, lgbm_probs, n_bins=10, strategy="uniform")
    ece = _expected_calibration_error(y, lgbm_probs, n_bins=10)
    slope, intercept = _calibration_slope_intercept(y, lgbm_probs)
    brier = float(np.mean((lgbm_probs - y) ** 2))

    calib_path = TABLES_DIR / "calibration_diagnostics.csv"
    with open(calib_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["metric", "value"])
        writer.writerow(["brier_score", round(brier, 4)])
        writer.writerow(["ece_10bin", round(ece, 4)])
        writer.writerow(["calibration_slope", round(slope, 4)])
        writer.writerow(["calibration_intercept", round(intercept, 4)])
        writer.writerow([])
        writer.writerow(["bin_mean_predicted", "bin_fraction_positive"])
        for pp, pt in zip(prob_pred, prob_true):
            writer.writerow([round(float(pp), 4), round(float(pt), 4)])

    print(f"\nCalibration diagnostics written: {calib_path}")
    print(f"  Brier score:            {brier:.4f}")
    print(f"  ECE (10-bin):           {ece:.4f}")
    print(f"  Calibration slope:      {slope:.4f} (1.0 = perfect)")
    print(f"  Calibration intercept:  {intercept:.4f} (0.0 = perfect)")


if __name__ == "__main__":
    run()