# Copyright 2026 The AiVa proof-of-concept team. Licensed under the Apache License, Version 2.0.
"""Score an RLCD checkpoint on a labelled split, the way the published model was judged.

    python evaluate.py --model rlcd-mine --data data/test.jsonl --latency 200 \
        --prior public=0.146,internal=0.527,confidential=0.297,restricted=0.024,sovereign-only=0.006

Reports, per head: accuracy, macro-F1, expected calibration error, and coverage / accuracy at confidence
thresholds (to choose your own). For the ordered sensitivity head, under the tighten-only rule (the highest
level L with P(class >= L) >= --tighten): how many sensitive requests are classified too low, overall and per
level, and how many public or internal ones are raised, with 95% Wilson intervals. If rows carry a "dlp"
level (your DLP's verdict), it also reports the rate with DLP alongside. `--prior` re-weights rows to your
production sensitivity mix: weighted figures are reported beside the as-sampled ones, never instead of them.
Macro-F1 averages over every class of a head, as the published figures do (a class absent from the split
scores 0); `--macro-f1-over present` averages over the classes present instead, for small or partial splits.
`--latency N` times N single requests on this machine, tokenizer included.
`--model` is a folder or a Hub repo; `--revision` pins a Hub revision.
"""

import argparse
import itertools
import json
import math
import time
from pathlib import Path

import torch
from huggingface_hub import snapshot_download

from rlcd import RlcdClassifier, render


def wilson(hits: float, n: float) -> tuple[float, float]:
    if n == 0:
        return (0.0, 1.0)
    p, z = hits / n, 1.96
    centre, half = p + z * z / (2 * n), z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
    return ((centre - half) / (1 + z * z / n), (centre + half) / (1 + z * z / n))


def tighten(distribution: dict[str, float], order: list[str], threshold: float) -> str:
    mass = 0.0
    for level in reversed(order):
        mass += distribution[level]
        if mass >= threshold:
            return level
    return order[0]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", required=True)
    parser.add_argument("--revision", default=None)
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--prior", default=None, help="level=share,... to weight rows to production")
    parser.add_argument("--tighten", type=float, default=0.5)
    parser.add_argument("--sensitive-from", default="confidential", help="the first level that counts as sensitive")
    parser.add_argument("--thresholds", default="0.5,0.6,0.7,0.8,0.9,0.95")
    parser.add_argument("--macro-f1-over", choices=["all", "present"], default="all")
    parser.add_argument("--latency", type=int, default=0)
    parser.add_argument("--out", type=Path, default=None, help="write the report as JSON")
    args = parser.parse_args()

    path = Path(args.model) if Path(args.model).is_dir() else Path(snapshot_download(args.model, revision=args.revision))
    model = RlcdClassifier(path)
    heads, temperatures = model.config["heads"], model.config["temperatures"]
    rows = [json.loads(line) for line in args.data.read_text().splitlines() if line.strip()]
    texts = [r.get("text") or render(r["messages"]) for r in rows]

    probs: dict[str, list[torch.Tensor]] = {h: [] for h in heads}
    with torch.no_grad():
        for i in range(0, len(texts), 16):
            logits = model.logits(texts[i : i + 16])
            for h in heads:
                probs[h].append((logits[h] / temperatures[h]).softmax(-1))
    probs = {h: torch.cat(p) for h, p in probs.items()}

    weights = torch.ones(len(rows))
    if args.prior:
        prior = {k: float(v) for k, v in (item.split("=") for item in args.prior.split(","))}
        levels = [r.get("sensitivity") for r in rows]
        share = {lvl: levels.count(lvl) / len(levels) for lvl in set(levels)}
        weights = torch.tensor([prior.get(lvl, 0.0) / share[lvl] for lvl in levels])

    report: dict = {"rows": len(rows), "provenance": "measured", "heads": {}}
    thresholds = [float(t) for t in args.thresholds.split(",")]

    def metrics(p: torch.Tensor, labels: torch.Tensor, w: torch.Tensor, classes: int) -> dict:
        confidence, prediction = p.max(-1)
        correct = prediction.eq(labels).float()
        over = range(classes) if args.macro_f1_over == "all" else sorted(set(labels.tolist()) | set(prediction.tolist()))
        f1s = []
        for c in over:
            tp = (w * ((prediction == c) & (labels == c))).sum()
            fp, fn = (w * ((prediction == c) & (labels != c))).sum(), (w * ((prediction != c) & (labels == c))).sum()
            f1s.append(0.0 if tp == 0 else (2 * tp / (2 * tp + fp + fn)).item())
        ece = 0.0
        for low, high in itertools.pairwise(torch.linspace(0, 1, 16)):
            in_bin = (confidence > low) & (confidence <= high)
            if in_bin.any():
                bw = w[in_bin]
                gap = ((confidence[in_bin] * bw).sum() - (correct[in_bin] * bw).sum()) / bw.sum()
                ece += (bw.sum() / w.sum()).item() * abs(gap.item())
        curve = {}
        for t in thresholds:
            kept = confidence >= t
            accuracy = ((correct[kept] * w[kept]).sum() / w[kept].sum()).item() if kept.any() else None
            curve[str(t)] = {"coverage": round((w[kept].sum() / w.sum()).item(), 4),
                             "accuracy": None if accuracy is None else round(accuracy, 4)}
        return {"accuracy": round(((correct * w).sum() / w.sum()).item(), 4),
                "macro_f1": round(sum(f1s) / max(1, len(f1s)), 4), "ece": round(ece, 4), "coverage_accuracy": curve}

    for h, classes in heads.items():
        keep = [i for i, r in enumerate(rows) if r.get(h) in classes]
        if not keep:
            continue
        labels = torch.tensor([classes.index(rows[i][h]) for i in keep])
        report["heads"][h] = {"rows": len(keep), **metrics(probs[h][keep], labels, torch.ones(len(keep)), len(classes))}
        if args.prior:  # beside the as-sampled figures, never instead of them
            report["heads"][h]["under_prior"] = metrics(probs[h][keep], labels, weights[keep], len(classes))

    if "sensitivity" in heads:
        order = heads["sensitivity"]
        start = order.index(args.sensitive_from)
        rank = order.index
        safety: dict = {"tighten_at": args.tighten}
        variants = {"rlcd": lambda i, lvl: lvl}
        if any("dlp" in r for r in rows):
            variants["rlcd_or_dlp"] = lambda i, lvl: max(lvl, rows[i].get("dlp") or order[0], key=rank)
        for name, combine in variants.items():
            effective = [combine(i, tighten(dict(zip(order, probs["sensitivity"][i].tolist(), strict=True)), order, args.tighten))
                         for i in range(len(rows))]
            sensitive = [i for i, r in enumerate(rows) if r.get("sensitivity") in order[start:]]
            under = [i for i in sensitive if rank(effective[i]) < rank(rows[i]["sensitivity"])]
            total, hits = sum(weights[i].item() for i in sensitive), sum(weights[i].item() for i in under)
            entry = {"under_classified": {"rate": round(len(under) / len(sensitive), 4) if sensitive else None,
                                          "of": len(sensitive), "ci95": [round(x, 4) for x in wilson(len(under), len(sensitive))],
                                          **({"rate_under_prior": round(hits / total, 4) if total else None} if args.prior else {})}}
            for level in order[start:]:
                at = [i for i in sensitive if rows[i]["sensitivity"] == level]
                n_under = sum(1 for i in at if rank(effective[i]) < rank(level))
                entry[f"under_{level}"] = {"rate": round(n_under / len(at), 4) if at else None, "of": len(at),
                                           "ci95": [round(x, 4) for x in wilson(n_under, len(at))]}
            plain = [i for i, r in enumerate(rows) if r.get("sensitivity") in order[:start]]
            for level in order[start:]:
                raised = sum(1 for i in plain if rank(effective[i]) >= rank(level))
                entry[f"raised_to_{level}_or_above"] = {"rate": round(raised / len(plain), 4) if plain else None,
                                                        "of": len(plain)}
            safety[name] = entry
        report["sensitivity_safety"] = safety

    if args.latency:
        sample = [rows[i % len(rows)] for i in range(args.latency)]
        timings = []
        for row in sample:
            started = time.perf_counter()
            model.predict(row.get("messages") or [{"role": "user", "content": row["text"]}])
            timings.append((time.perf_counter() - started) * 1000)
        timings.sort()
        report["latency_ms"] = {"p50": round(timings[len(timings) // 2], 1),
                                "p95": round(timings[min(len(timings) - 1, int(0.95 * len(timings)))], 1),
                                "requests": len(timings), "threads": torch.get_num_threads()}
    print(json.dumps(report, indent=1))
    if args.out:
        args.out.write_text(json.dumps(report, indent=1))


if __name__ == "__main__":
    main()