scripts/evaluate_head.py

"""
Evaluate CLIP head checkpoints on cached embedding splits.

This is the Stage 3A checkpoint-comparison script. It loads one candidate head,
optionally a Stage 2 baseline head, and reports overall, per-generator,
per-source, per-model-family, and per-augmentation metrics for any requested
embedding splits.

Usage
-----
    python scripts/evaluate_head.py \\
        --emb-dir data/embeddings \\
        --candidate data/checkpoints/head_v3a.pt \\
        --baseline data/checkpoints/head_stage2.pt \\
        --split test \\
        --split heldout \\
        --split test_augmented \\
        --report-out data/reports/head_v3a_eval.json
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np

from train_head import SplitEmbeddings, _build_head, _load_split


def _softmax(logits):
    import torch

    return torch.softmax(logits, dim=-1)


def _rate(num: int, den: int) -> float:
    return float(num / den) if den else 0.0


def _threshold_metrics(logits, labels, uncertainty_threshold: float) -> dict[str, Any]:
    probs = _softmax(logits)
    max_probs, argmax = probs.max(dim=-1)
    certain = max_probs >= uncertainty_threshold

    authentic = labels == 0
    ai = labels == 1
    pred_authentic = argmax == 0
    pred_ai = argmax == 1

    tp = int((certain & pred_ai & ai).sum().item())
    tn = int((certain & pred_authentic & authentic).sum().item())
    fp = int((certain & pred_ai & authentic).sum().item())
    fn = int((certain & pred_authentic & ai).sum().item())
    uncertain_authentic = int((~certain & authentic).sum().item())
    uncertain_ai = int((~certain & ai).sum().item())
    uncertain = uncertain_authentic + uncertain_ai
    n = int(len(labels))

    argmax_correct = int((argmax == labels).sum().item())
    certain_correct = int(((argmax == labels) & certain).sum().item())
    certain_n = int(certain.sum().item())

    return {
        "n": n,
        "argmax_accuracy": _rate(argmax_correct, n),
        "coverage_accuracy": _rate(certain_correct, certain_n),
        "coverage_rate": _rate(certain_n, n),
        "uncertainty_rate": _rate(uncertain, n),
        "false_positive_rate": _rate(fp, fp + tn + uncertain_authentic),
        "false_negative_rate": _rate(fn, fn + tp + uncertain_ai),
        "confusion": {
            "tp": tp,
            "tn": tn,
            "fp": fp,
            "fn": fn,
            "uncertain_authentic": uncertain_authentic,
            "uncertain_ai": uncertain_ai,
        },
        "mean_confidence": float(max_probs.mean().item()) if n else 0.0,
    }


def _indices_for(values: np.ndarray, label: str) -> list[int]:
    return [i for i, value in enumerate(values) if str(value) == label]


def _group_metrics(
    logits,
    labels,
    values: np.ndarray,
    uncertainty_threshold: float,
) -> dict[str, dict[str, Any]]:
    import torch

    result: dict[str, dict[str, Any]] = {}
    group_labels = sorted({str(v) for v in values if str(v)})
    for label in group_labels:
        idx = _indices_for(values, label)
        if not idx:
            continue
        tensor_idx = torch.as_tensor(idx, dtype=torch.long, device=logits.device)
        result[label] = _threshold_metrics(
            logits.index_select(0, tensor_idx),
            labels.index_select(0, tensor_idx),
            uncertainty_threshold,
        )
    return result


def _generator_values(split: SplitEmbeddings) -> np.ndarray:
    values: list[str] = []
    labels = split.y.cpu().numpy()
    for i, label in enumerate(labels):
        if label != 1:
            values.append("")
            continue
        values.append(str(split.generators[i] or split.sources[i]))
    return np.asarray(values)


def _augmentation_values(split: SplitEmbeddings) -> np.ndarray:
    return np.asarray(
        [str(value) if str(value) else "clean" for value in split.augmentations]
    )


def _split_report(
    split: SplitEmbeddings,
    logits,
    uncertainty_threshold: float,
) -> dict[str, Any]:
    report: dict[str, Any] = {
        "overall": _threshold_metrics(logits, split.y, uncertainty_threshold)
    }

    groups = {
        "by_source": split.sources,
        "by_generator": _generator_values(split),
        "by_model_family": split.model_families,
    }
    if any(str(value) for value in split.augmentations):
        groups["by_augmentation"] = _augmentation_values(split)

    for name, values in groups.items():
        metrics = _group_metrics(logits, split.y, values, uncertainty_threshold)
        if metrics:
            report[name] = metrics

    return report


def _load_head(checkpoint: Path):
    import torch

    head = _build_head()
    state = torch.load(checkpoint, map_location="cpu")
    head.load_state_dict(state)
    head.eval()
    return head


def _evaluate_checkpoint(
    checkpoint: Path,
    splits: dict[str, SplitEmbeddings],
    uncertainty_threshold: float,
) -> dict[str, Any]:
    import torch

    head = _load_head(checkpoint)
    report: dict[str, Any] = {
        "checkpoint": str(checkpoint),
        "splits": {},
    }
    with torch.no_grad():
        for name, split in splits.items():
            logits = head(split.x)
            report["splits"][name] = _split_report(
                split,
                logits,
                uncertainty_threshold,
            )
    return report


def _comparison(candidate: dict[str, Any], baseline: dict[str, Any] | None) -> dict:
    if baseline is None:
        return {}

    result: dict[str, dict[str, float]] = {}
    for split, candidate_report in candidate["splits"].items():
        if split not in baseline["splits"]:
            continue
        candidate_overall = candidate_report["overall"]
        baseline_overall = baseline["splits"][split]["overall"]
        result[split] = {
            "argmax_accuracy_delta": (
                candidate_overall["argmax_accuracy"]
                - baseline_overall["argmax_accuracy"]
            ),
            "uncertainty_rate_delta": (
                candidate_overall["uncertainty_rate"]
                - baseline_overall["uncertainty_rate"]
            ),
            "false_positive_rate_delta": (
                candidate_overall["false_positive_rate"]
                - baseline_overall["false_positive_rate"]
            ),
            "false_negative_rate_delta": (
                candidate_overall["false_negative_rate"]
                - baseline_overall["false_negative_rate"]
            ),
        }
    return result


def _print_summary(report: dict[str, Any]) -> None:
    print(f"Candidate: {report['candidate']['checkpoint']}")
    if report.get("baseline"):
        print(f"Baseline: {report['baseline']['checkpoint']}")
    for split, metrics in report["candidate"]["splits"].items():
        overall = metrics["overall"]
        print(
            f"  {split}: n={overall['n']} "
            f"acc={overall['argmax_accuracy']:.4f} "
            f"uncertain={overall['uncertainty_rate']:.4f} "
            f"fpr={overall['false_positive_rate']:.4f} "
            f"fnr={overall['false_negative_rate']:.4f}"
        )
        for group_name in ["by_generator", "by_augmentation"]:
            if group_name not in metrics:
                continue
            print(f"    {group_name}:")
            for label, group_metrics in metrics[group_name].items():
                print(
                    f"      {label}: n={group_metrics['n']} "
                    f"acc={group_metrics['argmax_accuracy']:.4f} "
                    f"uncertain={group_metrics['uncertainty_rate']:.4f}"
                )


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--emb-dir", type=Path, required=True)
    parser.add_argument("--candidate", type=Path, required=True)
    parser.add_argument("--baseline", type=Path, default=None)
    parser.add_argument(
        "--split",
        action="append",
        default=[],
        help="Embedding split name to evaluate, without .npz. Defaults to test.",
    )
    parser.add_argument("--report-out", type=Path, required=True)
    parser.add_argument("--uncertainty-threshold", type=float, default=0.6)
    args = parser.parse_args()

    split_names = args.split or ["test"]
    splits = {name: _load_split(args.emb_dir, name) for name in split_names}

    candidate = _evaluate_checkpoint(
        args.candidate,
        splits,
        args.uncertainty_threshold,
    )
    baseline = (
        _evaluate_checkpoint(args.baseline, splits, args.uncertainty_threshold)
        if args.baseline is not None
        else None
    )

    report = {
        "uncertainty_threshold": args.uncertainty_threshold,
        "candidate": candidate,
        "baseline": baseline,
        "comparison": _comparison(candidate, baseline),
    }

    args.report_out.parent.mkdir(parents=True, exist_ok=True)
    with args.report_out.open("w", encoding="utf-8") as fh:
        json.dump(report, fh, indent=2, sort_keys=True)

    _print_summary(report)
    print(f"\nSaved evaluation report to {args.report_out}")


if __name__ == "__main__":
    main()