"""Evaluate a trained run: top-1 vs jev_topic, macro-F1, per-item dump, confidence buckets.

Mirrors the eval/ schema promised in jev-topic-head (per_item jsonl + headline numbers).
No rules_table.json exists in this example, so show/hide is skipped (it's a fixed
lookup table downstream of the topic, not part of the head).
"""
import argparse, json, os

import numpy as np
import torch
import torch.nn.functional as F
import yaml
from sklearn.metrics import f1_score
from transformers import AutoModel, AutoTokenizer
from peft import PeftModel

from train_head import PrototypeHead, encode, read_jsonl


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="config.yaml")
    ap.add_argument("--run", required=True, help="e.g. runs/seed0")
    args = ap.parse_args()
    cfg = yaml.safe_load(open(args.config))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 if cfg["train"]["dtype"] == "bf16" else torch.float32

    hs = torch.load(os.path.join(args.run, "head", "head.pt"), map_location="cpu")
    topics = hs["topics"]
    topic2id = {t: i for i, t in enumerate(topics)}

    tok = AutoTokenizer.from_pretrained(cfg["base_model"])
    tok.padding_side = "right"
    base = AutoModel.from_pretrained(hs.get("base_model", cfg["base_model"]),
                                     torch_dtype=dtype).to(device)
    base.config.use_cache = False
    model = PeftModel.from_pretrained(base, os.path.join(args.run, "lora"))

    head = PrototypeHead(hs["protos"], float(hs["scale"]))
    head.load_state_dict({"protos": hs["protos"], "scale": hs["scale"]})
    head = head.to(device).eval()

    rows = read_jsonl(os.path.join(cfg["data"]["dir"], cfg["data"]["test"]))
    texts = [r["text"] for r in rows]
    logits, BATCH = [], 32
    with torch.no_grad():
        for i in range(0, len(texts), BATCH):
            v = encode(model, tok, texts[i:i + BATCH], cfg["max_length"], device)
            logits.append(head.logits(v).float().cpu())
    logits = torch.cat(logits)
    probs = F.softmax(logits, dim=-1)
    conf, pred = probs.max(-1)
    pred = pred.numpy()
    gold = np.array([topic2id[r["jev_topic"]] for r in rows])

    per_item = []
    for r, p, c in zip(rows, pred, conf.numpy()):
        per_item.append({"id": r["id"], "model_topic": topics[p],
                         "jev_topic": r["jev_topic"], "confidence": round(float(c), 4),
                         "agree": bool(p == topic2id[r["jev_topic"]])})
    os.makedirs(os.path.join(args.run, "eval"), exist_ok=True)
    with open(os.path.join(args.run, "eval", "per_item_test.jsonl"), "w") as f:
        for r in per_item:
            f.write(json.dumps(r) + "\n")

    top1 = float((pred == gold).mean())
    mf1 = float(f1_score(gold, pred, average="macro", zero_division=0,
                         labels=list(range(len(topics)))))
    conf_all = conf.numpy()
    summary = {"n": len(rows), "top1_vs_jev": round(top1, 4),
               "macro_f1_vs_jev": round(mf1, 4),
               "conf_lt_05": int((conf_all < 0.5).sum()),
               "conf_gte_09": int((conf_all >= 0.9).sum()),
               "errors_lt_05": int(((pred != gold) & (conf_all < 0.5)).sum()),
               "errors_gte_09": int(((pred != gold) & (conf_all >= 0.9)).sum())}
    json.dump(summary, open(os.path.join(args.run, "eval", "summary.json"), "w"), indent=2)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()