urlmodel_info/evaluation.py

"""
Evaluation utilities — metrics, comparison tables, and publication-quality plots.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    roc_auc_score,
    confusion_matrix,
    classification_report,
)


# ═════════════════════════════════════════════════════════════════════════════
# Core evaluation
# ═════════════════════════════════════════════════════════════════════════════

def evaluate_model(
    name: str,
    y_true: np.ndarray,
    y_pred: np.ndarray,
    y_proba: np.ndarray | None = None,
    class_names: list[str] | None = None,
    task: str = "binary",
) -> dict:
    """Compute comprehensive classification metrics.

    Parameters
    ----------
    name : str
        Model name.
    y_true, y_pred : array-like
        Ground truth and predicted labels.
    y_proba : array-like, optional
        Class probability matrix (n_samples × n_classes).
    class_names : list[str]
        Human-readable class names.
    task : str
        ``"binary"`` or ``"multiclass"``.

    Returns
    -------
    dict with keys: name, accuracy, precision, recall, f1_macro, auc_roc,
         confusion_matrix, classification_report.
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    acc = accuracy_score(y_true, y_pred)
    prec = precision_score(y_true, y_pred, average="macro", zero_division=0)
    rec = recall_score(y_true, y_pred, average="macro", zero_division=0)
    f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)

    auc = None
    if y_proba is not None:
        y_proba = np.asarray(y_proba)
        try:
            if task == "binary" and y_proba.ndim == 2:
                auc = roc_auc_score(y_true, y_proba[:, 1])
            else:
                auc = roc_auc_score(y_true, y_proba, multi_class="ovr", average="macro")
        except Exception:
            auc = None

    cm = confusion_matrix(y_true, y_pred)
    target_names = class_names or [str(c) for c in sorted(np.unique(y_true))]
    cr = classification_report(y_true, y_pred, target_names=target_names, zero_division=0)

    return {
        "name": name,
        "accuracy": acc,
        "precision": prec,
        "recall": rec,
        "f1_macro": f1,
        "auc_roc": auc,
        "confusion_matrix": cm,
        "classification_report": cr,
    }


def print_results_banner(result: dict) -> None:
    """Pretty-print a model's evaluation results."""
    print(f"\n{'─' * 50}")
    print(f"  {result['name']}")
    print(f"{'─' * 50}")
    print(f"  Accuracy  : {result['accuracy']:.4f}")
    print(f"  Precision : {result['precision']:.4f}")
    print(f"  Recall    : {result['recall']:.4f}")
    print(f"  F1 (macro): {result['f1_macro']:.4f}")
    if result["auc_roc"] is not None:
        print(f"  AUC-ROC   : {result['auc_roc']:.4f}")
    print(f"{'─' * 50}")


# ═════════════════════════════════════════════════════════════════════════════
# Literature comparison table
# ═════════════════════════════════════════════════════════════════════════════

_LITERATURE_BASELINES = [
    {"Method": "RF on URL features (Sahingoz 2019)", "Accuracy": 0.9798,
     "F1 (macro)": 0.9793, "AUC-ROC": None, "Source": "Literature"},
    {"Method": "XGBoost (Rao & Pais 2019)", "Accuracy": 0.9685,
     "F1 (macro)": None, "AUC-ROC": 0.993, "Source": "Literature"},
    {"Method": "URLNet CharCNN (Le 2018)", "Accuracy": 0.985,
     "F1 (macro)": None, "AUC-ROC": 0.986, "Source": "Literature"},
    {"Method": "LSTM (Bahnsen 2018)", "Accuracy": 0.987,
     "F1 (macro)": None, "AUC-ROC": None, "Source": "Literature"},
    {"Method": "Stacking CNN+LSTM+GRU (Nguyen 2021)", "Accuracy": 0.975,
     "F1 (macro)": 0.972, "AUC-ROC": None, "Source": "Literature"},
    {"Method": "Hybrid DNN+LSTM (Ozcan 2021)", "Accuracy": 0.964,
     "F1 (macro)": None, "AUC-ROC": 0.973, "Source": "Literature"},
    {"Method": "CharCNN (Aljofey 2022)", "Accuracy": 0.9675,
     "F1 (macro)": 0.9671, "AUC-ROC": 0.991, "Source": "Literature"},
]


def build_comparison_table(
    our_results: list[dict],
    task: str = "binary",
) -> pd.DataFrame:
    """Build a comparison table of our results vs. literature baselines.

    Parameters
    ----------
    our_results : list[dict]
        Each dict is the output of ``evaluate_model()``.
    task : str
        ``"binary"`` or ``"multiclass"``.

    Returns
    -------
    pd.DataFrame
    """
    rows = list(_LITERATURE_BASELINES)
    for r in our_results:
        rows.append({
            "Method": r["name"],
            "Accuracy": r["accuracy"],
            "F1 (macro)": r["f1_macro"],
            "AUC-ROC": r["auc_roc"],
            "Source": "Ours",
        })
    df = pd.DataFrame(rows)
    df = df.sort_values("Accuracy", ascending=False).reset_index(drop=True)
    return df


# ═════════════════════════════════════════════════════════════════════════════
# Plotting helpers
# ═════════════════════════════════════════════════════════════════════════════

def plot_accuracy_bar(
    comparison_df: pd.DataFrame,
    save_path: str | None = None,
) -> plt.Figure:
    """Horizontal bar chart of accuracies, colour-coded by source."""
    fig, ax = plt.subplots(figsize=(10, max(5, len(comparison_df) * 0.45)))
    colors = ["#1976D2" if s == "Ours" else "#BDBDBD"
              for s in comparison_df["Source"]]
    y = range(len(comparison_df))
    ax.barh(y, comparison_df["Accuracy"], color=colors, edgecolor="white", height=0.6)
    ax.set_yticks(y)
    ax.set_yticklabels(comparison_df["Method"], fontsize=9)
    ax.set_xlabel("Accuracy")
    ax.set_title("Model Accuracy — Ours vs Literature")
    ax.invert_yaxis()
    ax.set_xlim(0.85, 1.01)
    for i, v in enumerate(comparison_df["Accuracy"]):
        if pd.notna(v):
            ax.text(v + 0.002, i, f"{v:.4f}", va="center", fontsize=8)
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig


def plot_confusion_matrix(
    cm: np.ndarray,
    class_names: list[str],
    title: str = "Confusion Matrix",
    save_path: str | None = None,
) -> plt.Figure:
    """Plot an annotated confusion matrix heatmap."""
    fig, ax = plt.subplots(figsize=(6, 5))
    sns.heatmap(
        cm, annot=True, fmt="d", cmap="Blues",
        xticklabels=class_names, yticklabels=class_names, ax=ax,
    )
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title(title)
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig


def plot_feature_importance(
    feature_names: list[str],
    importances: np.ndarray,
    top_n: int = 25,
    title: str = "Feature Importance",
    save_path: str | None = None,
) -> plt.Figure:
    """Horizontal bar chart of top-N feature importances."""
    idx = np.argsort(importances)[-top_n:]
    fig, ax = plt.subplots(figsize=(8, max(4, top_n * 0.3)))
    ax.barh(
        [feature_names[i] for i in idx],
        importances[idx],
        color="#1976D2", edgecolor="white",
    )
    ax.set_xlabel("Importance")
    ax.set_title(title)
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    return fig