#!/usr/bin/env python3
"""Element-head evaluation on held-out sites: does each image get a rendering Haiku judged acceptable?

Compares the v2 element head with what v1 shipped for the same images (its 5-class image head for the first 32
images of a page, the hand-written runtime rules for everything else), per kind of image.
"""
import argparse, collections, gzip, importlib.util, json, pickle, sys
from pathlib import Path
import numpy as np
import torch

from model import NadirNet, TREAT
from train import collate, prep, element_logits
import safety

V1_MAP = {"keep": 0, "dim": 0, "invert": 1, "light": 1, "plate": 1, "clear": 3}


def v1_decision(t, im, surf_L):
    """v1 runtime: model treatment (or 'keep' for images it never saw) through the hand-written rules."""
    s = safety.safe_treatment(t, im, surf_L)
    if s.startswith("adapt:"):
        return 2  # the logo adaptation path
    return V1_MAP.get(s, 0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--crawl", action="append", required=True, help="TAG=CRAWL_DIR")
    ap.add_argument("--v1-ckpt", default="")
    ap.add_argument("--v1-src", default="", help="v1 model.py (its architecture)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--rows", default="", help="also write every element (probabilities, labels) as jsonl")
    args = ap.parse_args()
    crawls = dict(c.split("=", 1) for c in args.crawl)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    ck = torch.load(args.ckpt, map_location="cpu")
    T = float(ck.get("temperature", 1.0))
    m = NadirNet(**ck["cfg"]); m.load_state_dict(ck["model"]); m.eval().to(dev)
    v1 = None
    if args.v1_ckpt and args.v1_src:
        spec = importlib.util.spec_from_file_location("model_v1", args.v1_src)
        mv1 = importlib.util.module_from_spec(spec); spec.loader.exec_module(mv1)
        c1 = torch.load(args.v1_ckpt, map_location="cpu")
        v1 = mv1.NadirNet(**c1["cfg"]); v1.load_state_dict(c1["model"]); v1.eval().to(dev)
    data = pickle.load(open(args.data, "rb"))
    val = [prep(s) for s in data["val"] if len(s.get("E_ut", []))]
    rows = []
    with torch.no_grad():
        for s in val:
            b = collate([s], dev)
            lab, _, ctx = m(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
            tl, rl = element_logits(m, b, lab.float(), ctx)
            k = int(b["e_mask"][0].sum())
            p = torch.softmax(tl[0, :k] / T, -1).cpu().numpy()
            role = rl[0, :k].argmax(-1).cpu().numpy()
            rprob = torch.softmax(rl[0, :k].float() / T, -1).max(-1).values.cpu().numpy()
            v1t = None
            if v1 is not None and len(s["IF"]):
                _, _, il = v1(b["tok_f"], b["role"], b["lab"], b["img_f"][..., :39], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
                v1t = il[0].float().argmax(-1).cpu().numpy()
            L = json.load(gzip.open(Path(crawls[s["crawl"]]) / "pages" / s["id"] / "light.json.gz"))
            n = len(s["role"])
            outL = lab[0, :n, 0].float().cpu().numpy()
            for q in range(k):
                j = int(s["E_idx"][q])
                im = L["images"][j]
                u = im.get("underTok", -1)
                surf_L = float(outL[u]) if 0 <= u < n else None
                ud = [round(float(x), 4) for x in lab[0, u].float().cpu().numpy()] if 0 <= u < n else None
                old = ["keep", "dim", "invert", "light", "plate"][int(v1t[j])] if (v1t is not None and j < 32 and j < len(v1t)) else "keep"
                soft = s["E_soft"][q]
                rows.append({"bucket": s["E_bucket"][q], "kind": s["E_kind"][q], "soft": soft.tolist(), "best": int(np.argmax(soft)),
                             "v2": int(np.argmax(p[q])), "conf": float(p[q].max()), "p": [round(float(x), 4) for x in p[q]], "role_pred": int(role[q]), "role": int(s["E_role"][q]),
                             "v1": v1_decision(old, im, surf_L), "v1_seen": bool(v1t is not None and j < 32), "nlab": int(s["E_nlab"][q]),
                             **({"rp": round(float(rprob[q]), 4), "f": [round(float(x), 4) for x in s["E_f"][q]], "ud": ud, "ul": [round(float(x), 4) for x in s["E_ul"][q]]} if args.rows else {})})
    def score(rs, key):
        if not rs:
            return {}
        acc = np.mean([r[key] == r["best"] for r in rs])
        ok = np.mean([r["soft"][r[key]] > 0 for r in rs])
        return {"n": len(rs), "acc": round(float(acc), 4), "acceptable": round(float(ok), 4)}
    rep = {"temperature": T, "overall": {"v2": score(rows, "v2"), "v1": score(rows, "v1")}, "by_bucket": {}, "by_kind": {}}
    for key in ("bucket", "kind"):
        groups = collections.defaultdict(list)
        for r in rows:
            groups[r[key]].append(r)
        rep["by_" + key] = {g: {"v2": score(rs, "v2"), "v1": score(rs, "v1")} for g, rs in sorted(groups.items())}
    # calibration: confidence vs accuracy
    bins = collections.defaultdict(list)
    for r in rows:
        bins[min(9, int(r["conf"] * 10))].append(r["v2"] == r["best"])
    rep["calibration"] = {f"{b / 10:.1f}": {"n": len(v), "acc": round(float(np.mean(v)), 3)} for b, v in sorted(bins.items())}
    rr = [r for r in rows if r["role"] >= 0]
    rep["role_acc"] = round(float(np.mean([r["role_pred"] == r["role"] for r in rr])), 4) if rr else None
    # Haiku's own consistency where a second, independent label exists
    rep["n_double_labelled"] = sum(r["nlab"] > 1 for r in rows)
    json.dump(rep, open(args.out, "w"), indent=1)
    if args.rows:
        with open(args.rows, "w") as f:
            for r in rows:
                f.write(json.dumps(r) + "\n")
    print(json.dumps({k: rep[k] for k in ("overall", "role_acc", "temperature")}, indent=1))
    print(json.dumps(rep["by_bucket"], indent=1))


if __name__ == "__main__":
    main()