"""Holdout evaluation for trained DualHeadModernBERT models.

Loads a trained checkpoint, runs inference on the 1,200-paragraph holdout set,
and compares predictions against proxy gold labels (GPT-5.4, Opus-4.6) or
real human gold labels when available.

Reports: macro/per-class F1, precision, recall, MCC, AUC (one-vs-rest),
QWK, MAE, Krippendorff's alpha, confusion matrices, calibration (ECE),
cost/latency comparison, and generates publication-ready figures.
"""

import json
import time
from dataclasses import dataclass
from pathlib import Path

import krippendorff
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
import torch.nn.functional as F
from safetensors.torch import load_file
from sklearn.metrics import (
    f1_score,
    precision_score,
    recall_score,
    matthews_corrcoef,
    roc_auc_score,
    cohen_kappa_score,
    confusion_matrix,
    classification_report,
    mean_absolute_error,
)
from transformers import AutoModel, AutoTokenizer

from .data import CATEGORIES, CAT2ID, NUM_CATEGORIES, NUM_SPECIFICITY
from .model import DualHeadModernBERT, ordinal_predict, softmax_predict

SPEC_LABELS = ["L1: Generic", "L2: Domain", "L3: Firm-Specific", "L4: Quantified"]


@dataclass
class EvalConfig:
    """Configuration for holdout evaluation."""
    checkpoint_path: str
    paragraphs_path: str
    holdout_path: str
    benchmark_paths: dict[str, str]  # name → path to benchmark JSONL
    output_dir: str
    max_seq_length: int = 512
    batch_size: int = 64
    # Architecture params (must match training)
    specificity_head: str = "independent"
    spec_mlp_dim: int = 256
    pooling: str = "attention"


def _load_jsonl(path: str | Path) -> list[dict]:
    records = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def load_holdout_data(
    paragraphs_path: str,
    holdout_path: str,
    benchmark_paths: dict[str, str],
) -> dict:
    """Load holdout paragraphs and benchmark labels, joined by paragraph ID."""
    paragraphs = {p["id"]: p for p in _load_jsonl(paragraphs_path)}

    with open(holdout_path) as f:
        holdout_ids = json.load(f)

    # Load benchmark labels
    benchmarks = {}
    for name, path in benchmark_paths.items():
        labels = {r["paragraphId"]: r for r in _load_jsonl(path)}
        benchmarks[name] = labels

    # Build joined dataset
    records = []
    for pid in holdout_ids:
        para = paragraphs.get(pid)
        if para is None:
            continue

        bench_labels = {}
        for name, labels in benchmarks.items():
            if pid in labels:
                lab = labels[pid]["label"]
                bench_labels[name] = {
                    "category": lab["content_category"],
                    "specificity": lab["specificity_level"],
                }

        records.append({
            "id": pid,
            "text": para["text"],
            "benchmark_labels": bench_labels,
        })

    return records


def load_model(config: EvalConfig, device: torch.device) -> tuple:
    """Load trained DualHeadModernBERT from checkpoint."""
    checkpoint = Path(config.checkpoint_path)

    tokenizer = AutoTokenizer.from_pretrained(str(checkpoint))

    # Load backbone (we need the architecture, then overwrite weights)
    try:
        import flash_attn  # noqa: F401
        attn_impl = "flash_attention_2"
    except ImportError:
        attn_impl = "sdpa"

    backbone = AutoModel.from_pretrained(
        "answerdotai/ModernBERT-large",
        trust_remote_code=True,
        attn_implementation=attn_impl,
        dtype=torch.bfloat16,
    )
    hidden_size = backbone.config.hidden_size

    model = DualHeadModernBERT(
        backbone=backbone,
        hidden_size=hidden_size,
        num_categories=NUM_CATEGORIES,
        num_specificity=NUM_SPECIFICITY,
        specificity_head_type=config.specificity_head,
        spec_mlp_dim=config.spec_mlp_dim,
        pooling=config.pooling,
    )

    # Load trained weights (filter out loss function buffers not needed for inference)
    state_dict = load_file(str(checkpoint / "model.safetensors"))
    model.load_state_dict(state_dict, strict=False)
    model = model.to(device).to(torch.bfloat16)
    model.eval()

    return model, tokenizer


@torch.no_grad()
def run_inference(
    model: DualHeadModernBERT,
    tokenizer,
    records: list[dict],
    max_seq_length: int,
    batch_size: int,
    device: torch.device,
    spec_head_type: str = "independent",
) -> dict:
    """Run model inference on holdout records. Returns predictions + timing."""
    texts = [r["text"] for r in records]
    all_cat_logits = []
    all_spec_logits = []
    total_time = 0.0

    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i : i + batch_size]
        encoded = tokenizer(
            batch_texts,
            truncation=True,
            max_length=max_seq_length,
            padding="longest",
            return_tensors="pt",
        ).to(device)

        start = time.perf_counter()
        outputs = model(
            input_ids=encoded["input_ids"],
            attention_mask=encoded["attention_mask"],
        )
        torch.cuda.synchronize() if device.type == "cuda" else None
        total_time += time.perf_counter() - start

        all_cat_logits.append(outputs["category_logits"].float().cpu())
        all_spec_logits.append(outputs["specificity_logits"].float().cpu())

    cat_logits = torch.cat(all_cat_logits, dim=0)
    spec_logits = torch.cat(all_spec_logits, dim=0)

    cat_probs = F.softmax(cat_logits, dim=1).numpy()
    cat_preds = cat_logits.argmax(dim=1).numpy()

    if spec_head_type == "softmax":
        spec_preds = softmax_predict(spec_logits).numpy()
        spec_probs = F.softmax(spec_logits, dim=1).numpy()
    else:
        spec_preds = ordinal_predict(spec_logits).numpy()
        # Convert ordinal logits to class probs for AUC
        spec_probs = _ordinal_to_class_probs(spec_logits).numpy()

    return {
        "cat_preds": cat_preds,
        "cat_probs": cat_probs,
        "cat_logits": cat_logits.numpy(),
        "spec_preds": spec_preds,
        "spec_probs": spec_probs,
        "spec_logits": spec_logits.numpy(),
        "total_time_s": total_time,
        "num_samples": len(texts),
        "avg_ms_per_sample": (total_time / len(texts)) * 1000,
    }


def _ordinal_to_class_probs(logits: torch.Tensor) -> torch.Tensor:
    """Convert ordinal threshold logits to per-class probabilities.

    P(class=0) = 1 - P(>=1)
    P(class=k) = P(>=k) - P(>=k+1) for 0 < k < K-1
    P(class=K-1) = P(>=K-1)
    """
    probs = torch.sigmoid(logits)  # (N, K-1)
    num_classes = probs.shape[1] + 1
    class_probs = torch.zeros(probs.shape[0], num_classes)

    class_probs[:, 0] = 1.0 - probs[:, 0]
    for k in range(1, num_classes - 1):
        class_probs[:, k] = probs[:, k - 1] - probs[:, k]
    class_probs[:, -1] = probs[:, -1]

    # Clamp for numerical stability
    class_probs = class_probs.clamp(min=0.0)
    class_probs = class_probs / class_probs.sum(dim=1, keepdim=True)

    return class_probs


def compute_ece(probs: np.ndarray, labels: np.ndarray, n_bins: int = 15) -> tuple[float, dict]:
    """Compute Expected Calibration Error and per-bin data for reliability diagram."""
    confidences = np.max(probs, axis=1)
    predictions = np.argmax(probs, axis=1)
    accuracies = (predictions == labels).astype(float)

    bin_edges = np.linspace(0.0, 1.0, n_bins + 1)
    bin_data = {"bin_centers": [], "bin_accs": [], "bin_confs": [], "bin_counts": []}
    ece = 0.0

    for i in range(n_bins):
        mask = (confidences > bin_edges[i]) & (confidences <= bin_edges[i + 1])
        count = mask.sum()
        if count > 0:
            avg_conf = confidences[mask].mean()
            avg_acc = accuracies[mask].mean()
            ece += (count / len(labels)) * abs(avg_acc - avg_conf)
            bin_data["bin_centers"].append((bin_edges[i] + bin_edges[i + 1]) / 2)
            bin_data["bin_accs"].append(avg_acc)
            bin_data["bin_confs"].append(avg_conf)
            bin_data["bin_counts"].append(int(count))

    return ece, bin_data


def compute_all_metrics(
    preds: np.ndarray,
    labels: np.ndarray,
    probs: np.ndarray,
    label_names: list[str],
    task_name: str,
    is_ordinal: bool = False,
) -> dict:
    """Compute comprehensive metrics for a single classification task."""
    num_classes = len(label_names)

    # Basic classification metrics
    macro_f1 = f1_score(labels, preds, average="macro", labels=range(num_classes))
    weighted_f1 = f1_score(labels, preds, average="weighted", labels=range(num_classes))
    per_class_f1 = f1_score(labels, preds, average=None, labels=range(num_classes))
    per_class_precision = precision_score(labels, preds, average=None, labels=range(num_classes), zero_division=0)
    per_class_recall = recall_score(labels, preds, average=None, labels=range(num_classes), zero_division=0)
    macro_precision = precision_score(labels, preds, average="macro", labels=range(num_classes), zero_division=0)
    macro_recall = recall_score(labels, preds, average="macro", labels=range(num_classes), zero_division=0)

    # MCC (multiclass)
    mcc = matthews_corrcoef(labels, preds)

    # AUC (one-vs-rest, macro)
    try:
        auc = roc_auc_score(labels, probs, multi_class="ovr", average="macro", labels=range(num_classes))
    except ValueError:
        auc = float("nan")

    # Confusion matrix
    cm = confusion_matrix(labels, preds, labels=range(num_classes))

    # Calibration
    ece, bin_data = compute_ece(probs, labels)

    metrics = {
        f"{task_name}_macro_f1": macro_f1,
        f"{task_name}_weighted_f1": weighted_f1,
        f"{task_name}_macro_precision": macro_precision,
        f"{task_name}_macro_recall": macro_recall,
        f"{task_name}_mcc": mcc,
        f"{task_name}_auc": auc,
        f"{task_name}_ece": ece,
        f"{task_name}_confusion_matrix": cm.tolist(),
        f"{task_name}_calibration": bin_data,
    }

    for i, name in enumerate(label_names):
        short = name.replace(" ", "").replace("/", "").replace(":", "")[:8]
        metrics[f"{task_name}_f1_{short}"] = per_class_f1[i]
        metrics[f"{task_name}_prec_{short}"] = per_class_precision[i]
        metrics[f"{task_name}_recall_{short}"] = per_class_recall[i]

    # Ordinal-specific metrics
    if is_ordinal:
        metrics[f"{task_name}_qwk"] = cohen_kappa_score(labels, preds, weights="quadratic")
        metrics[f"{task_name}_mae"] = mean_absolute_error(labels, preds)

    # Krippendorff's alpha (treat model + reference as two coders)
    reliability_data = np.array([labels, preds])
    if is_ordinal:
        metrics[f"{task_name}_kripp_alpha"] = krippendorff.alpha(
            reliability_data=reliability_data, level_of_measurement="ordinal"
        )
    else:
        metrics[f"{task_name}_kripp_alpha"] = krippendorff.alpha(
            reliability_data=reliability_data, level_of_measurement="nominal"
        )

    return metrics


def generate_figures(
    results: dict,
    output_dir: Path,
    model_name: str,
    ref_name: str,
) -> list[str]:
    """Generate publication-ready figures. Returns list of saved file paths."""
    figures_dir = output_dir / "figures"
    figures_dir.mkdir(parents=True, exist_ok=True)
    saved = []

    sns.set_theme(style="whitegrid", font_scale=1.1)
    palette = sns.color_palette("viridis", 7)

    # 1. Category confusion matrix
    cm = np.array(results["cat_confusion_matrix"])
    fig, ax = plt.subplots(figsize=(10, 8))
    cm_norm = cm.astype(float) / cm.sum(axis=1, keepdims=True)
    sns.heatmap(
        cm_norm, annot=cm, fmt="d", cmap="Blues",
        xticklabels=[c[:12] for c in CATEGORIES],
        yticklabels=[c[:12] for c in CATEGORIES],
        ax=ax, vmin=0, vmax=1, cbar_kws={"label": "Proportion"},
    )
    ax.set_xlabel(f"Predicted ({model_name})")
    ax.set_ylabel(f"Reference ({ref_name})")
    ax.set_title(f"Category Confusion Matrix — {model_name} vs {ref_name}")
    plt.tight_layout()
    path = figures_dir / f"confusion_cat_{ref_name.lower().replace(' ', '_')}.png"
    fig.savefig(path, dpi=150)
    plt.close(fig)
    saved.append(str(path))

    # 2. Specificity confusion matrix
    cm_spec = np.array(results["spec_confusion_matrix"])
    fig, ax = plt.subplots(figsize=(7, 6))
    cm_spec_norm = cm_spec.astype(float) / cm_spec.sum(axis=1, keepdims=True).clip(min=1)
    sns.heatmap(
        cm_spec_norm, annot=cm_spec, fmt="d", cmap="Oranges",
        xticklabels=SPEC_LABELS,
        yticklabels=SPEC_LABELS,
        ax=ax, vmin=0, vmax=1, cbar_kws={"label": "Proportion"},
    )
    ax.set_xlabel(f"Predicted ({model_name})")
    ax.set_ylabel(f"Reference ({ref_name})")
    ax.set_title(f"Specificity Confusion Matrix — {model_name} vs {ref_name}")
    plt.tight_layout()
    path = figures_dir / f"confusion_spec_{ref_name.lower().replace(' ', '_')}.png"
    fig.savefig(path, dpi=150)
    plt.close(fig)
    saved.append(str(path))

    # 3. Calibration reliability diagram (category)
    cal = results["cat_calibration"]
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.bar(cal["bin_centers"], cal["bin_accs"], width=0.06, alpha=0.7, label="Accuracy", color="#4C72B0")
    ax.plot([0, 1], [0, 1], "k--", alpha=0.5, label="Perfect calibration")
    ax.set_xlabel("Confidence")
    ax.set_ylabel("Accuracy")
    ax.set_title(f"Category Calibration — ECE={results['cat_ece']:.4f}")
    ax.legend()
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    plt.tight_layout()
    path = figures_dir / f"calibration_cat_{ref_name.lower().replace(' ', '_')}.png"
    fig.savefig(path, dpi=150)
    plt.close(fig)
    saved.append(str(path))

    # 4. Per-class F1 bar chart
    cat_f1s = [results.get(f"cat_f1_{c.replace(' ', '').replace('/', '')[:8]}", 0) for c in CATEGORIES]
    spec_f1s = [results.get(f"spec_f1_{s.replace(' ', '').replace(':', '')[:8]}", 0) for s in SPEC_LABELS]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    bars = ax1.barh([c[:15] for c in CATEGORIES], cat_f1s, color=palette)
    ax1.axvline(0.80, color="red", linestyle="--", alpha=0.7, label="Target (0.80)")
    ax1.set_xlabel("F1 Score")
    ax1.set_title(f"Per-Category F1 ({model_name} vs {ref_name})")
    ax1.set_xlim(0, 1)
    ax1.legend()
    for bar, v in zip(bars, cat_f1s):
        ax1.text(v + 0.01, bar.get_y() + bar.get_height() / 2, f"{v:.3f}", va="center", fontsize=9)

    bars2 = ax2.barh(SPEC_LABELS, spec_f1s, color=sns.color_palette("YlOrRd", 4))
    ax2.axvline(0.80, color="red", linestyle="--", alpha=0.7, label="Target (0.80)")
    ax2.set_xlabel("F1 Score")
    ax2.set_title(f"Per-Level Specificity F1 ({model_name} vs {ref_name})")
    ax2.set_xlim(0, 1)
    ax2.legend()
    for bar, v in zip(bars2, spec_f1s):
        ax2.text(v + 0.01, bar.get_y() + bar.get_height() / 2, f"{v:.3f}", va="center", fontsize=9)

    plt.tight_layout()
    path = figures_dir / f"per_class_f1_{ref_name.lower().replace(' ', '_')}.png"
    fig.savefig(path, dpi=150)
    plt.close(fig)
    saved.append(str(path))

    return saved


def generate_comparison_figures(
    all_results: dict[str, dict],
    output_dir: Path,
) -> list[str]:
    """Generate figures comparing multiple models/references."""
    figures_dir = output_dir / "figures"
    figures_dir.mkdir(parents=True, exist_ok=True)
    saved = []

    sns.set_theme(style="whitegrid", font_scale=1.1)

    # 1. Model comparison bar chart (if we have CORAL vs Independent)
    model_names = list(all_results.keys())
    if len(model_names) >= 2:
        metric_keys = ["cat_macro_f1", "spec_macro_f1", "cat_mcc", "spec_qwk", "spec_mae"]
        metric_labels = ["Cat Macro F1", "Spec Macro F1", "Cat MCC", "Spec QWK", "Spec MAE"]

        fig, ax = plt.subplots(figsize=(12, 6))
        x = np.arange(len(metric_labels))
        width = 0.8 / len(model_names)

        for i, name in enumerate(model_names):
            vals = []
            for k in metric_keys:
                v = all_results[name].get(k, 0)
                vals.append(v if not np.isnan(v) else 0)
            bars = ax.bar(x + i * width - 0.4 + width / 2, vals, width, label=name, alpha=0.85)
            for bar, v in zip(bars, vals):
                ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                        f"{v:.3f}", ha="center", va="bottom", fontsize=8)

        ax.set_xticks(x)
        ax.set_xticklabels(metric_labels)
        ax.set_ylabel("Score")
        ax.set_title("Model Comparison — All Metrics")
        ax.legend()
        ax.axhline(0.80, color="red", linestyle="--", alpha=0.5, label="F1 Target")
        plt.tight_layout()
        path = figures_dir / "model_comparison.png"
        fig.savefig(path, dpi=150)
        plt.close(fig)
        saved.append(str(path))

    # 2. Speed/cost comparison
    speed_data = {}
    for name, res in all_results.items():
        if "avg_ms_per_sample" in res:
            speed_data[name] = res["avg_ms_per_sample"]

    if speed_data:
        # Add GenAI panel reference speeds from docs
        genai_speeds = {
            "GPT-5.4 (API)": 2900,  # from provenance data
            "Opus-4.6 (API)": 6000,  # from provenance data
        }

        fig, ax = plt.subplots(figsize=(10, 5))
        all_names = list(speed_data.keys()) + list(genai_speeds.keys())
        all_speeds = list(speed_data.values()) + list(genai_speeds.values())
        colors = ["#4C72B0"] * len(speed_data) + ["#DD8452"] * len(genai_speeds)

        bars = ax.barh(all_names, all_speeds, color=colors)
        ax.set_xlabel("Latency (ms per paragraph)")
        ax.set_title("Inference Speed Comparison")
        ax.set_xscale("log")
        for bar, v in zip(bars, all_speeds):
            ax.text(v * 1.1, bar.get_y() + bar.get_height() / 2,
                    f"{v:.1f}ms", va="center", fontsize=9)

        plt.tight_layout()
        path = figures_dir / "speed_comparison.png"
        fig.savefig(path, dpi=150)
        plt.close(fig)
        saved.append(str(path))

    return saved


def format_report(
    model_name: str,
    ref_name: str,
    metrics: dict,
    inference_stats: dict,
) -> str:
    """Format a human-readable evaluation report."""
    lines = []
    lines.append(f"\n{'='*70}")
    lines.append(f"  HOLDOUT EVALUATION: {model_name} vs {ref_name}")
    lines.append(f"{'='*70}\n")

    lines.append(f"  Samples evaluated: {inference_stats['num_samples']}")
    lines.append(f"  Total inference time: {inference_stats['total_time_s']:.2f}s")
    lines.append(f"  Avg latency: {inference_stats['avg_ms_per_sample']:.2f}ms/sample")
    lines.append(f"  Throughput: {1000 / inference_stats['avg_ms_per_sample']:.0f} samples/sec\n")

    # Category metrics
    lines.append(f"  {'─'*50}")
    lines.append(f"  CATEGORY CLASSIFICATION")
    lines.append(f"  {'─'*50}")
    lines.append(f"  Macro F1:       {metrics['cat_macro_f1']:.4f}  {'✓' if metrics['cat_macro_f1'] >= 0.80 else '✗'} (target: 0.80)")
    lines.append(f"  Weighted F1:    {metrics['cat_weighted_f1']:.4f}")
    lines.append(f"  Macro Prec:     {metrics['cat_macro_precision']:.4f}")
    lines.append(f"  Macro Recall:   {metrics['cat_macro_recall']:.4f}")
    lines.append(f"  MCC:            {metrics['cat_mcc']:.4f}")
    lines.append(f"  AUC (OvR):      {metrics['cat_auc']:.4f}")
    lines.append(f"  ECE:            {metrics['cat_ece']:.4f}")
    lines.append(f"  Kripp Alpha:    {metrics['cat_kripp_alpha']:.4f}")
    lines.append("")

    lines.append(f"  {'Category':<25} {'F1':>8} {'Prec':>8} {'Recall':>8}")
    lines.append(f"  {'-'*25} {'-'*8} {'-'*8} {'-'*8}")
    for c in CATEGORIES:
        short = c.replace(" ", "").replace("/", "")[:8]
        f1 = metrics.get(f"cat_f1_{short}", 0)
        prec = metrics.get(f"cat_prec_{short}", 0)
        rec = metrics.get(f"cat_recall_{short}", 0)
        lines.append(f"  {c:<25} {f1:>8.4f} {prec:>8.4f} {rec:>8.4f}")

    # Specificity metrics
    lines.append(f"\n  {'─'*50}")
    lines.append(f"  SPECIFICITY CLASSIFICATION")
    lines.append(f"  {'─'*50}")
    lines.append(f"  Macro F1:       {metrics['spec_macro_f1']:.4f}  {'✓' if metrics['spec_macro_f1'] >= 0.80 else '✗'} (target: 0.80)")
    lines.append(f"  Weighted F1:    {metrics['spec_weighted_f1']:.4f}")
    lines.append(f"  Macro Prec:     {metrics['spec_macro_precision']:.4f}")
    lines.append(f"  Macro Recall:   {metrics['spec_macro_recall']:.4f}")
    lines.append(f"  MCC:            {metrics['spec_mcc']:.4f}")
    lines.append(f"  AUC (OvR):      {metrics['spec_auc']:.4f}")
    lines.append(f"  QWK:            {metrics['spec_qwk']:.4f}")
    lines.append(f"  MAE:            {metrics['spec_mae']:.4f}")
    lines.append(f"  ECE:            {metrics['spec_ece']:.4f}")
    lines.append(f"  Kripp Alpha:    {metrics['spec_kripp_alpha']:.4f}")
    lines.append("")

    lines.append(f"  {'Level':<25} {'F1':>8} {'Prec':>8} {'Recall':>8}")
    lines.append(f"  {'-'*25} {'-'*8} {'-'*8} {'-'*8}")
    for s in SPEC_LABELS:
        short = s.replace(" ", "").replace(":", "")[:8]
        f1 = metrics.get(f"spec_f1_{short}", 0)
        prec = metrics.get(f"spec_prec_{short}", 0)
        rec = metrics.get(f"spec_recall_{short}", 0)
        lines.append(f"  {s:<25} {f1:>8.4f} {prec:>8.4f} {rec:>8.4f}")

    lines.append(f"\n{'='*70}\n")

    return "\n".join(lines)


def evaluate(config: EvalConfig) -> dict:
    """Run full holdout evaluation. Returns all metrics + generates figures."""
    output_dir = Path(config.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n  Device: {device}")

    # Load holdout data
    print("  Loading holdout data...")
    records = load_holdout_data(
        config.paragraphs_path,
        config.holdout_path,
        config.benchmark_paths,
    )
    print(f"  Holdout paragraphs: {len(records)}")

    # Load model
    print(f"  Loading model from {config.checkpoint_path}...")
    model, tokenizer = load_model(config, device)
    param_count = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"  Model parameters: {param_count:.0f}M")

    # Run inference
    print("  Running inference...")
    inference = run_inference(
        model, tokenizer, records,
        config.max_seq_length, config.batch_size,
        device, config.specificity_head,
    )
    print(f"  Inference complete: {inference['total_time_s']:.2f}s ({inference['avg_ms_per_sample']:.2f}ms/sample)")

    # Evaluate against each benchmark
    all_results = {}
    all_figures = []

    model_name = Path(config.checkpoint_path).parent.name
    if model_name == "final":
        model_name = Path(config.checkpoint_path).parent.parent.name

    for ref_name, ref_path in config.benchmark_paths.items():
        print(f"\n  Evaluating against {ref_name}...")

        # Get reference labels for paragraphs that have them
        cat_labels = []
        spec_labels = []
        cat_preds = []
        spec_preds = []
        cat_probs = []
        spec_probs = []

        for i, rec in enumerate(records):
            bench = rec["benchmark_labels"].get(ref_name)
            if bench is None:
                continue
            cat_labels.append(CAT2ID[bench["category"]])
            spec_labels.append(bench["specificity"] - 1)  # 0-indexed
            cat_preds.append(inference["cat_preds"][i])
            spec_preds.append(inference["spec_preds"][i])
            cat_probs.append(inference["cat_probs"][i])
            spec_probs.append(inference["spec_probs"][i])

        cat_labels = np.array(cat_labels)
        spec_labels = np.array(spec_labels)
        cat_preds = np.array(cat_preds)
        spec_preds = np.array(spec_preds)
        cat_probs = np.array(cat_probs)
        spec_probs = np.array(spec_probs)

        print(f"  Matched samples: {len(cat_labels)}")

        # Compute metrics
        cat_metrics = compute_all_metrics(
            cat_preds, cat_labels, cat_probs, CATEGORIES, "cat", is_ordinal=False
        )
        spec_metrics = compute_all_metrics(
            spec_preds, spec_labels, spec_probs, SPEC_LABELS, "spec", is_ordinal=True
        )

        combined = {**cat_metrics, **spec_metrics, **inference}
        combined["combined_macro_f1"] = (combined["cat_macro_f1"] + combined["spec_macro_f1"]) / 2

        # Print report
        report = format_report(model_name, ref_name, combined, inference)
        print(report)

        # Save report
        report_path = output_dir / f"report_{ref_name.lower().replace(' ', '_').replace('.', '')}.txt"
        with open(report_path, "w") as f:
            f.write(report)

        # Generate figures
        figs = generate_figures(combined, output_dir, model_name, ref_name)
        all_figures.extend(figs)

        all_results[f"{model_name}_vs_{ref_name}"] = combined

    # Cross-model comparison figures
    comp_figs = generate_comparison_figures(all_results, output_dir)
    all_figures.extend(comp_figs)

    # Save all metrics as JSON
    serializable = {}
    for k, v in all_results.items():
        serializable[k] = {
            mk: mv for mk, mv in v.items()
            if isinstance(mv, (int, float, str, list, bool))
        }
    metrics_path = output_dir / "metrics.json"
    with open(metrics_path, "w") as f:
        json.dump(serializable, f, indent=2, default=str)

    print(f"\n  Results saved to {output_dir}")
    print(f"  Figures: {len(all_figures)} generated")
    for fig_path in all_figures:
        print(f"    {fig_path}")

    return all_results


def evaluate_comparison(
    configs: list[EvalConfig],
    output_dir: str,
) -> dict:
    """Evaluate multiple model checkpoints and generate cross-model comparisons."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    all_results = {}
    for config in configs:
        results = evaluate(config)
        all_results.update(results)

    # Generate comparison figures across all models
    comp_figs = generate_comparison_figures(all_results, output_dir)
    print(f"\n  Comparison figures: {len(comp_figs)}")

    return all_results
