"""Metrics from eval_student.py prediction dumps, per set and per variant (all / no settings_destructive), plus gaps.

Per set: per-question accuracy vs gold, record-exact (all non-box questions right), teacher agreement (argmax),
teacher accuracy (reference), NLL, ECE(15) before and after temperature scaling, box mean IoU and IoU>=0.5 per box type
(gold-present boxes only), box-confidence ECE vs IoU>=0.5, presence accuracy per box type (presence predicted from the
paired question in meta.box_pairs: primary_button != none, other pairs == "true"; gold = gold box present), and
record-exact per scenario.
Temperature scaling: one global T (grid search on NLL). "ece_ts_xfit" = 2-fold cross-fit inside the set (folds by
seeded id hash; T fitted on one half, applied to the other); "ece_ts_valT" = T fitted on --val-set (if given) applied here.

usage: python scripts/student_report.py --set gallery_eval=runs/p1.jsonl --set style_eval=runs/p2.jsonl ...
           [--val-set gallery_eval] [--gap gallery_train_vs_eval=gallery_train300,gallery_eval ...] --json out.json
"""
import argparse, hashlib, json, math, sys
from collections import Counter, defaultdict

SD = "settings_destructive"


def load(p):
    with open(p, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


def softmax(z, T=1.0):
    m = max(v / T for v in z); e = [math.exp(v / T - m) for v in z]; s = sum(e)
    return [v / s for v in e]


def ece(confs, hits, bins=15):
    n = len(confs); tot = 0.0
    if not n:
        return None
    for b in range(bins):
        lo, hi = b / bins, (b + 1) / bins
        idx = [i for i, c in enumerate(confs) if lo < c <= hi or (b == 0 and c == 0)]
        if idx:
            tot += len(idx) / n * abs(sum(hits[i] for i in idx) / len(idx) - sum(confs[i] for i in idx) / len(idx))
    return tot


def items(recs):
    """(opts, logits, gold_index) for every scored non-box question."""
    out = []
    for r in recs:
        for q, lz in r["logits"].items():
            g = r["gold"].get(q)
            if isinstance(g, dict) or g is None or str(g) not in lz:
                continue
            opts = list(lz); out.append((opts, [lz[o] for o in opts], opts.index(str(g))))
    return out


def nll_at(its, T):
    return sum(-math.log(max(softmax(z, T)[y], 1e-12)) for _, z, y in its) / max(len(its), 1)


GRID = [math.exp(math.log(0.2) + i * (math.log(5.0) - math.log(0.2)) / 240) for i in range(241)]


def fit_T(its):
    return min(GRID, key=lambda T: nll_at(its, T)) if its else 1.0


def ece_at(its, T):
    confs, hits = [], []
    for _, z, y in its:
        p = softmax(z, T); k = max(range(len(p)), key=p.__getitem__); confs.append(p[k]); hits.append(int(k == y))
    return ece(confs, hits)


def fold(rid, seed=0):
    return int(hashlib.md5(f"{seed}:{rid}".encode()).hexdigest(), 16) % 2


def metrics(recs, val_T=None):
    if not recs:
        return {"n_records": 0}
    perq = defaultdict(lambda: [0, 0]); agree = defaultdict(lambda: [0, 0]); tacc = defaultdict(lambda: [0, 0])
    exact = 0; n_exact = 0; confs, hits = [], []; scen = defaultdict(lambda: [0, 0]); conf_mat = defaultdict(Counter)
    for r in recs:
        ok_all, any_q = True, False
        for q, p in r["probs"].items():
            g = r["gold"].get(q); pred = max(p, key=p.get)
            t = r["teacher"].get(q)
            if t:
                tp = max(t, key=t.get); agree[q][0] += int(tp == pred); agree[q][1] += 1
            if g is None or isinstance(g, dict):
                continue
            any_q = True; hit = int(pred == str(g)); perq[q][0] += hit; perq[q][1] += 1
            confs.append(p[pred]); hits.append(hit); ok_all &= bool(hit)
            if not hit:
                conf_mat[q][f"{g}->{pred}"] += 1
            if t:
                tacc[q][0] += int(max(t, key=t.get) == str(g)); tacc[q][1] += 1
        if any_q:
            n_exact += 1; exact += int(ok_all)
            s = (r.get("meta") or {}).get("scenario") or r["id"]; scen[s][0] += int(ok_all); scen[s][1] += 1
    its = items(recs); f0 = [x for r in recs if fold(r["id"]) == 0 for x in items([r])]
    f1 = [x for r in recs if fold(r["id"]) == 1 for x in items([r])]
    T0, T1 = fit_T(f0), fit_T(f1)
    xc, xh = [], []
    for part, T in ((f1, T0), (f0, T1)):
        for _, z, y in part:
            p = softmax(z, T); k = max(range(len(p)), key=p.__getitem__); xc.append(p[k]); xh.append(int(k == y))
    # boxes
    iou = defaultdict(list); bconf = defaultdict(list); pres = defaultdict(lambda: [0, 0])
    for r in recs:
        pairs = (r.get("meta") or {}).get("box_pairs") or {}
        for bq, b in r.get("boxes", {}).items():
            if b.get("iou") is not None:
                iou[bq].append(b["iou"]); bconf[bq].append((b["conf"], int(b["iou"] >= 0.5)))
            g = r["gold"].get(bq); pq = pairs.get(bq)
            if isinstance(g, dict) and pq and pq in r["probs"]:
                p = r["probs"][pq]; pred = max(p, key=p.get)
                pp = (pred != "none") if pq == "primary_button" else (pred == "true")
                pres[bq][0] += int(pp == bool(g.get("present", True))); pres[bq][1] += 1
    allious = [v for vs in iou.values() for v in vs]; allc = [c for vs in bconf.values() for c in vs]
    box = None
    if iou or pres:
        box = {"overall": {"n": len(allious), "mean_iou": sum(allious) / len(allious) if allious else None,
                           "iou50": sum(v >= 0.5 for v in allious) / len(allious) if allious else None,
                           "iou75": sum(v >= 0.75 for v in allious) / len(allious) if allious else None,
                           "conf_ece_iou50": ece([c for c, _ in allc], [h for _, h in allc]),
                           "presence_acc": (sum(v[0] for v in pres.values()) / sum(v[1] for v in pres.values())) if pres else None,
                           "presence_n": sum(v[1] for v in pres.values())},
               "per_type": {bq: {"n": len(iou.get(bq, [])),
                                 "mean_iou": (sum(iou[bq]) / len(iou[bq])) if iou.get(bq) else None,
                                 "iou50": (sum(v >= 0.5 for v in iou[bq]) / len(iou[bq])) if iou.get(bq) else None,
                                 "presence_acc": (pres[bq][0] / pres[bq][1]) if pres[bq][1] else None,
                                 "presence_n": pres[bq][1]} for bq in sorted(set(iou) | set(pres))}}
    n = max(len(hits), 1); Tall = fit_T(its)
    out = {"n_records": len(recs), "n_questions": len(hits), "accuracy": sum(hits) / n,
           "record_exact": exact / max(n_exact, 1),
           "teacher_agreement": (sum(v[0] for v in agree.values()) / sum(v[1] for v in agree.values())) if agree else None,
           "teacher_accuracy": (sum(v[0] for v in tacc.values()) / sum(v[1] for v in tacc.values())) if tacc else None,
           "nll": nll_at(its, 1.0), "ece_15bin": ece(confs, hits),
           "ece_ts_xfit": ece(xc, xh), "T_xfit": [T0, T1], "T_fit_on_this_set": Tall, "ece_ts_in_sample": ece_at(its, Tall),
           "per_question": {q: {"acc": v[0] / v[1], "n": v[1],
                                "teacher_agreement": (agree[q][0] / agree[q][1]) if agree[q][1] else None,
                                "teacher_acc": (tacc[q][0] / tacc[q][1]) if tacc[q][1] else None,
                                "top_errors": conf_mat[q].most_common(4)} for q, v in sorted(perq.items())},
           "box": box,
           "record_exact_per_scenario": {s: [v[0], v[1]] for s, v in sorted(scen.items())},
           "mean_sec": sum(r.get("sec", 0) for r in recs) / len(recs)}
    if val_T is not None:
        out["ece_ts_valT"] = ece_at(its, val_T); out["valT"] = val_T
    return out


KEYS = ["accuracy", "record_exact", "teacher_agreement", "nll", "ece_15bin", "ece_ts_xfit"]


def flat(m):
    d = {k: m.get(k) for k in KEYS}
    b = (m.get("box") or {}).get("overall") or {}
    d.update({"box_mean_iou": b.get("mean_iou"), "box_iou50": b.get("iou50"), "presence_acc": b.get("presence_acc")})
    return d


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--set", action="append", default=[], help="name=preds.jsonl")
    ap.add_argument("--val-set", default=None, help="set name whose fitted T is also applied to every set")
    ap.add_argument("--gap", action="append", default=[], help="name=setA,setB  (A minus B, both variants)")
    ap.add_argument("--json", default=None)
    a = ap.parse_args()
    sets = {k: load(v) for k, v in (s.split("=", 1) for s in a.set)}
    valT = fit_T(items(sets[a.val_set])) if a.val_set else None
    res = {"sets": {}, "gaps": {}, "val_set": a.val_set, "valT": valT}
    for name, recs in sets.items():
        res["sets"][name] = {"all": metrics(recs, valT), "no_settings_destructive": metrics([r for r in recs if SD not in r["id"]], valT)}
    for g in a.gap:
        name, pair = g.split("=", 1); A, B = pair.split(",")
        res["gaps"][name] = {}
        for var in ("all", "no_settings_destructive"):
            fa, fb = flat(res["sets"][A][var]), flat(res["sets"][B][var])
            res["gaps"][name][var] = {k: (fa[k] - fb[k]) if fa[k] is not None and fb[k] is not None else None for k in fa}
    summary = {n: {v: flat(m[v]) | {"n": m[v]["n_records"]} for v in m} for n, m in res["sets"].items()}
    res["summary"] = summary
    sys.stdout.reconfigure(encoding="utf-8")
    print(json.dumps({"summary": summary, "gaps": res["gaps"]}, indent=1))
    if a.json:
        with open(a.json, "w", encoding="utf-8", newline="\n") as f:
            json.dump(res, f, indent=1, ensure_ascii=False)


if __name__ == "__main__":
    main()