"""Teacher-vs-gold breakdown for GlintGallery-style JSONL (meta.scenario / meta.theme / ...).
Per question: accuracy, NLL, ECE. Per <group> x question: accuracy, most common wrong prediction, example ids.
Usage: python scripts/teacher_breakdown.py data/gallery-v1/train.clef.jsonl data/gallery-v1/eval.clef.jsonl
       [--group scenario] [--by theme,appearance,lang,content_size] [--worst 15] [--json out.json]"""
import argparse, json, math, sys
from collections import Counter, defaultdict
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from glint.data import load_jsonl, normalize_teacher
from scripts.teacher_report import ece


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("files", nargs="+")
    ap.add_argument("--group", default="scenario"); ap.add_argument("--by", default="theme,appearance,lang,content_size")
    ap.add_argument("--worst", type=int, default=15); ap.add_argument("--examples", type=int, default=4)
    ap.add_argument("--json", default=None); args = ap.parse_args()
    by = [b for b in args.by.split(",") if b]
    q_stats = defaultdict(lambda: {"n": 0, "hit": 0, "nll": 0.0, "confs": [], "hits": []})
    cell = defaultdict(lambda: {"n": 0, "hit": 0, "wrong": Counter(), "ex": [], "gold": None, "p_gold": 0.0})
    facet = {b: defaultdict(lambda: [0, 0]) for b in by}
    for fn in args.files:
        for ex in load_jsonl(fn):
            t, g, m = normalize_teacher(ex), ex.get("gold") or {}, ex.get("meta") or {}
            for qid, gv in g.items():
                d = t.get(qid)
                if d is None:
                    continue
                gv = str(gv); pred = max(d, key=d.get); h = int(pred == gv)
                s = q_stats[qid]; s["n"] += 1; s["hit"] += h; s["nll"] -= math.log(max(d.get(gv, 0.0), 1e-9))
                s["confs"].append(d[pred]); s["hits"].append(h)
                c = cell[(m.get(args.group, "?"), qid)]; c["n"] += 1; c["hit"] += h; c["gold"] = gv; c["p_gold"] += d.get(gv, 0.0)
                if not h:
                    c["wrong"][pred] += 1
                    if len(c["ex"]) < args.examples:
                        c["ex"].append({"id": ex["id"], "image": str(Path(ex["_base"]) / ex["images"][0]), "pred": pred,
                                        "p_pred": round(d[pred], 3), "p_gold": round(d.get(gv, 0.0), 3)})
                for b in by:
                    facet[b][(m.get(b, "?"), qid)][0] += h; facet[b][(m.get(b, "?"), qid)][1] += 1
    allc = [x for s in q_stats.values() for x in s["confs"]]; allh = [x for s in q_stats.values() for x in s["hits"]]
    n = max(len(allh), 1)
    rep = {"overall": {"n": len(allh), "accuracy": sum(allh) / n, "nll": sum(s["nll"] for s in q_stats.values()) / n,
                       "ece_15bin": ece(allc, allh)},
           "per_question": {q: {"n": s["n"], "accuracy": s["hit"] / s["n"], "nll": s["nll"] / s["n"],
                                "ece_15bin": ece(s["confs"], s["hits"])} for q, s in q_stats.items()}}
    cells = [{"group": k[0], "q": k[1], "n": c["n"], "accuracy": c["hit"] / c["n"], "gold": c["gold"],
              "mean_p_gold": c["p_gold"] / c["n"], "wrong_preds": dict(c["wrong"].most_common(3)), "examples": c["ex"]}
             for k, c in cell.items()]
    cells.sort(key=lambda x: (x["accuracy"], -x["n"]))
    rep["worst_cells"] = [c for c in cells if c["accuracy"] < 1.0][: args.worst]
    rep["n_imperfect_cells"] = sum(c["accuracy"] < 1.0 for c in cells)
    for b in by:
        rep[f"by_{b}"] = {}
        for (v, q), (h, k) in sorted(facet[b].items()):
            rep[f"by_{b}"].setdefault(v, {})[q] = round(h / k, 4)
    txt = json.dumps(rep, indent=1, ensure_ascii=False)
    print(txt)
    if args.json:
        Path(args.json).write_text(txt, encoding="utf-8")

if __name__ == "__main__":
    main()