src/models/figures.py

"""Evaluation figures for the model report and dashboard."""
from __future__ import annotations

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import precision_recall_curve, roc_curve

from src import config

PRIMARY = "#1f4e79"
ACCENT = "#c0392b"


def make_figures(y_true, prob, pred, threshold) -> None:
    config.FIGURES_DIR.mkdir(parents=True, exist_ok=True)

    # PR curve
    prec, rec, _ = precision_recall_curve(y_true, prob)
    fig, ax = plt.subplots(figsize=(5, 4))
    ax.plot(rec, prec, color=PRIMARY, lw=2)
    ax.set_xlabel("Recall"); ax.set_ylabel("Precision"); ax.set_title("Precision–Recall")
    ax.grid(alpha=0.3); fig.tight_layout()
    fig.savefig(config.FIGURES_DIR / "pr_curve.png", dpi=130); plt.close(fig)

    # ROC curve
    fpr, tpr, _ = roc_curve(y_true, prob)
    fig, ax = plt.subplots(figsize=(5, 4))
    ax.plot(fpr, tpr, color=PRIMARY, lw=2)
    ax.plot([0, 1], [0, 1], "--", color="gray", lw=1)
    ax.set_xlabel("False Positive Rate"); ax.set_ylabel("True Positive Rate"); ax.set_title("ROC")
    ax.grid(alpha=0.3); fig.tight_layout()
    fig.savefig(config.FIGURES_DIR / "roc_curve.png", dpi=130); plt.close(fig)

    # Confusion matrix
    from sklearn.metrics import confusion_matrix
    cm = confusion_matrix(y_true, pred)
    fig, ax = plt.subplots(figsize=(4, 3.5))
    im = ax.imshow(cm, cmap="Blues")
    for (i, j), v in np.ndenumerate(cm):
        ax.text(j, i, str(v), ha="center", va="center",
                color="white" if v > cm.max() / 2 else "black", fontsize=12)
    ax.set_xticks([0, 1]); ax.set_yticks([0, 1])
    ax.set_xticklabels(["Legit", "Mule"]); ax.set_yticklabels(["Legit", "Mule"])
    ax.set_xlabel("Predicted"); ax.set_ylabel("Actual"); ax.set_title("Confusion (test)")
    fig.tight_layout(); fig.savefig(config.FIGURES_DIR / "confusion_matrix.png", dpi=130); plt.close(fig)

    # Risk score distribution (threshold-anchored: 50 == decision threshold)
    score = np.array([config.prob_to_risk(p, threshold) for p in prob])
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(score[y_true == 0], bins=40, alpha=0.6, label="Legit", color=PRIMARY, density=True)
    ax.hist(score[y_true == 1], bins=40, alpha=0.7, label="Mule", color=ACCENT, density=True)
    ax.axvline(50, color="black", ls="--", lw=1.5, label="Decision threshold (50)")
    ax.set_xlabel("Risk score (0–100)"); ax.set_ylabel("Density")
    ax.set_title("Risk score distribution by class"); ax.legend()
    fig.tight_layout(); fig.savefig(config.FIGURES_DIR / "risk_distribution.png", dpi=130); plt.close(fig)