"""© KAND CA 2026 - evaluate a router checkpoint without retraining.

Reuses train_router_head.py's own loaders and report() so a published model and
a new one are scored by identical code in one session.

A checkpoint saved by train_router_head.py is a directory with the backbone plus
router_model.pt (the span scorer). The published Nawah-Router-v3 has the same
layout, so both load through the same path.
"""
import os, sys
os.environ.setdefault("DATA_DIR", "ds/data")
import torch
from transformers import AutoTokenizer
import train_router_head as T

def load_model(mid):
    m = T.RouterModel(mid)
    w = os.path.join(mid, "router_model.pt") if os.path.isdir(mid) else None
    if w is None:
        from huggingface_hub import hf_hub_download
        w = hf_hub_download(mid, "router_model.pt")
    if os.path.exists(w):
        m.load_state_dict(torch.load(w, map_location="cpu", weights_only=True))
    return m.cuda().eval()

for mid in sys.argv[1:]:
    tok = AutoTokenizer.from_pretrained(mid)
    m = load_model(mid)
    n = sum(p.numel() for p in m.parameters())
    print(f"\n{'='*66}\n{mid}   params={n:,}")
    for name in ("eval_unseen_lanes", "eval_unseen_domain", "eval_unseen_axis", "eval_hard"):
        rows = T.load(name)
        T.report(m, tok, rows, name.replace("eval_", ""))
    del m; torch.cuda.empty_cache()