"""Score the model on labeled forensic probes.

Verdict scoring: normalized keyword overlap between the model's constrained
verdict and the probe's expected answer. Format scoring: did it produce a
verdict and a confidence at all.

Usage:
  .venv/bin/python research/eval.py --ckpt ckpt/distill
"""

import argparse
import json
import re
import sys
from pathlib import Path

import torch

from model.config import TinyLiquidConfig, CONFIGS
from model.tiny_liquid import TinyLiquid
from model.utils import latest_ckpt
from data.tokenizer import load_tokenizer
from research.structured import analyst_report
from research import eval_labels as EL

WORDS = re.compile(r"[a-z]+")


def norm(s: str):
    return set(w for w in WORDS.findall(s.lower()) if len(w) > 2)


def verdict_score(got: str, expected: str) -> float:
    g, e = norm(got), norm(expected)
    if not e:
        return 0.0
    return len(g & e) / len(e)


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ckpt/distill")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--probes", default="data/eval_probes.jsonl")
    ap.add_argument("--max-scratch", type=int, default=90)
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--resume-from", default=None,
                    help="log file: skip probe ids already scored in the last eval section")
    return ap.parse_args()


def main():
    args = parse_args()
    torch.set_num_threads(args.threads)
    tok = load_tokenizer(args.tok)
    from pathlib import Path
    ckpt_path = Path(args.ckpt)
    if ckpt_path.is_file():
        ckpt = ckpt_path
    else:
        ckpt = latest_ckpt(args.ckpt)
    # Read zip archive into buffer for torch.load
    import io
    with open(ckpt, 'rb') as f:
        ckpt_data = f.read()
    sd = torch.load(io.BytesIO(ckpt_data), map_location="cpu", weights_only=False)
    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
                           **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
    cfg.mtp_heads = 0  # MTP pretrain-only; eval builds without MTP heads
    model = TinyLiquid(cfg)
    model.load_state_dict(sd["model"], strict=False)
    model.eval()
    print(f"== eval {ckpt} ==\n", flush=True)

    raw = [json.loads(l) for l in Path(args.probes).read_text().splitlines() if l.strip()]
    probes = []
    for i, p in enumerate(raw):
        if "expected" in p:
            probes.append({"id": p.get("id", "p%02d" % i), "cat": p.get("task", "generic"),
                           "persona": p.get("persona", "analyst"),
                           "user": p["user"], "expected": p["expected"]})
        else:
            probes.append({"id": p.get("task", "t") + "-%02d" % i,
                           "cat": p.get("task", "generic"),
                           "persona": "analyst", "user": p["user"],
                           "expected": p.get("expect", "")})
    if args.resume_from:
        from research.eval_summary import scored_ids
        done = scored_ids(args.resume_from)
        before = len(probes)
        probes = [p for p in probes if p["id"] not in done]
        if probes:
            print(f"[resume] skipping {before - len(probes)}/{before} already-scored probes; "
                  f"remaining {len(probes)}", flush=True)
    scores, formats, quals = [], 0, 0
    cat_acc = {}
    rows = []
    for p in probes:
        persona_id = 2 if p["persona"] == "skeptic" else 1
        r = analyst_report(model, tok, p["user"], persona_id=persona_id,
                           max_scratch=args.max_scratch)
        canon = EL.CANON.get(p["id"])
        qual = canon is None
        ok_format = bool(r["verdict"]) and bool(r["confidence"])
        formats += int(ok_format)
        if qual:
            sc = float("nan")
            quals += 1
        else:
            sc = 1.0 if r["verdict"].strip().lower() == canon else 0.0
            scores.append(sc)
            cat_acc.setdefault(p["cat"], []).append(sc)
        rows.append((p["id"], sc, r["verdict"], r["confidence"]))
        print(f"[{p['id']}] {'qual' if qual else '%.2f' % sc} | verdict: {r['verdict']} | conf: {r['confidence']}",
              flush=True)
    n = len(scores)
    acc = sum(scores) / n if n else float("nan")
    print(f"\ncanonical verdict accuracy (exact): {acc:.3f}  "
          f"(n={n}, qualitative={quals})  format rate: {formats/len(probes):.2f}")
    print("by category:")
    for cat, v in sorted(cat_acc.items()):
        print(f"  {cat:12s} acc {sum(v)/len(v):.3f}  n={len(v)}")


if __name__ == "__main__":
    main()