"""Gate a trained router on a held-out test set.

Usage:
  python evaluate_router.py --model router-model --test my-sealed-test.jsonl \
      --baseline-correct 172 --repeats 3

The test file has the same format as the training data ({"text", "label"}).
Write it yourself, BEFORE you train, and never train on it.

It prints overall accuracy, a confusion matrix, the gate lines that matter for
a router, and latency, then exits non-zero if any predeclared line fails.

Latency is local PyTorch inference on whatever device runs this script. Set
--max-p95-ms for YOUR target hardware; it is not comparable to a TensorRT
engine measured client-side on a Jetson.
"""
import argparse, json, statistics, sys, time
from collections import Counter

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MEDIA = {"image", "audio", "video"}   # lanes that start expensive work


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--test", required=True)
    ap.add_argument("--repeats", type=int, default=1)
    ap.add_argument("--baseline-correct", type=int, help="correct calls by the router you are replacing, same cases")
    ap.add_argument("--max-chat-to-media", type=int, default=0)
    ap.add_argument("--max-media-to-chat", type=int, default=0)
    ap.add_argument("--max-p95-ms", type=float, default=100.0)
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(args.model)
    tok.truncation_side = "left"
    model = AutoModelForSequenceClassification.from_pretrained(args.model).to(device).eval()
    id2label = model.config.id2label

    rows = [json.loads(line) for line in open(args.test) if line.strip()]
    confusion = Counter()
    latencies = []
    with torch.no_grad():
        model(**tok("warm up", return_tensors="pt").to(device))          # first call pays setup cost
        for _ in range(args.repeats):
            for row in rows:
                t0 = time.perf_counter()
                enc = tok(row["text"], truncation=True, max_length=512, return_tensors="pt").to(device)
                pred = id2label[int(model(**enc).logits.argmax(-1))]
                if device == "cuda":
                    torch.cuda.synchronize()
                latencies.append((time.perf_counter() - t0) * 1000)
                confusion[(row["label"], pred)] += 1

    calls = sum(confusion.values())
    correct = sum(n for (gold, pred), n in confusion.items() if gold == pred)
    chat_to_media = sum(n for (gold, pred), n in confusion.items() if gold == "chat" and pred in MEDIA)
    media_to_chat = sum(n for (gold, pred), n in confusion.items() if gold in MEDIA | {"search"} and pred == "chat")
    lat = sorted(latencies)
    p50 = statistics.median(lat)
    p95 = lat[min(len(lat) - 1, int(0.95 * len(lat)))]

    labels = sorted({g for g, _ in confusion} | {p for _, p in confusion})
    print(f"accuracy: {correct}/{calls}")
    print("confusion (rows = gold, cols = predicted):")
    print("        " + " ".join(f"{l[:6]:>6}" for l in labels))
    for g in labels:
        print(f"{g[:6]:>6}  " + " ".join(f"{confusion[(g, p)]:>6}" for p in labels))
    print(f"latency ms: p50 {p50:.1f}  p95 {p95:.1f}  ({device})")

    gates = [
        ("chat sent to a media lane", chat_to_media <= args.max_chat_to_media, f"{chat_to_media} (max {args.max_chat_to_media})"),
        ("specialist request sent to chat", media_to_chat <= args.max_media_to_chat, f"{media_to_chat} (max {args.max_media_to_chat})"),
        ("p95 latency", p95 <= args.max_p95_ms, f"{p95:.1f} ms (max {args.max_p95_ms})"),
    ]
    if args.baseline_correct is not None:
        gates.append(("at least as accurate as the baseline", correct >= args.baseline_correct,
                      f"{correct} vs {args.baseline_correct}"))
    failed = False
    for name, ok, detail in gates:
        print(f"{'PASS' if ok else 'FAIL'}  {name}: {detail}")
        failed |= not ok
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()