#!/usr/bin/env python3
"""Every measured model on the gold frames none of them trained on, with the noise made visible.

    python tools/compare_unseen.py            # table, runs/clean_gold_comparison.json

Two sets of frames:

* the **core**: the 149 frames that were unseen by every model on 2026-10-02
  (``data/gold/unseen_core.json``). It never changes, so the progression of the models over time
  stays comparable;
* **all unseen** frames: the core plus every gold frame added since that no dataset has used.

Where the reviewer kept a proposed outline, the gold holds the pre-labelling model's own mask,
which overlaps itself perfectly. With ``data/gold/kept_machine.json`` (tools/gold_kept.py) such
outlines are ignored object by object for every model; without it, a frame a model pre-labelled
(``data/gold/prelabel.json``) leaves that model's measure altogether.
  It grows with the gold set and is the one to judge a candidate on.

A model is listed when its predictions on the current gold set exist under
``runs/_measure_<version>/gold_predictions`` (``train_run.py`` keeps them for a new run; after
the gold set grew, re-measure the older versions or they are compared on fewer frames).
For every model the script also reports, against the active model, the paired-bootstrap
interval of the difference in F1 and mask score and the share of draws it wins (tools/gate_stats.py).
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Dict, List, Optional

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import gate_stats  # noqa: E402

NAMES = ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
BEAMLINE_PREFIX = "yolo26n-seg-overlap-false"   # the model trained before the loop


def class_report(c) -> Dict[str, float]:
    tp, fp, fn, iou = c
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
    return {"tp": int(tp), "fp": int(fp), "fn": int(fn), "precision": precision, "recall": recall, "f1": f1,
            "mean_mask_iou": iou / tp if tp else 0.0, "n_matched": int(tp)}


def report(frames: Dict[str, gate_stats.FrameCounts]) -> Dict:
    rows = list(frames.values())
    per_class = {}
    for cid, name in enumerate(NAMES):
        c = gate_stats._total(rows, cid)
        if c[0] or c[1] or c[2]:
            per_class[name] = class_report(c)
    return {"per_class": per_class, "overall": class_report(gate_stats._total(rows, None)),
            "dataset": {"gold_images": len(rows)}}


def mask_score(o: Dict) -> float:
    return o["mean_mask_iou"] * (o["tp"] / (o["tp"] + o["fn"]) if (o["tp"] + o["fn"]) else 0.0)


def table(title: str, reps: Dict[str, Dict]) -> None:
    print(title)
    print(f"{'run':46} {'frames':>6} {'F1':>6} {'recall':>7} {'IoU matched':>12} {'mask score':>11}   loop/pin/crystal/face/ice F1")
    for k, r in reps.items():
        o, pc = r["overall"], r["per_class"]
        print(f"{k:46} {r['dataset']['gold_images']:6d} {o['f1']:6.3f} {o['recall']:7.2f} {o['mean_mask_iou']:12.3f} {mask_score(o):11.3f}   "
              + "/".join(f"{pc.get(c, {}).get('f1', 0):.2f}" for c in NAMES[:5]))


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=str(HERE.parent), help="AareLC_train checkout (default: this one)")
    ap.add_argument("--draws", type=int, default=2000)
    args = ap.parse_args(argv)
    root = Path(args.root).expanduser()
    sys.path.insert(0, str(root))
    gold = root / "data" / "gold" / "labels"
    core_file = root / "data" / "gold" / "unseen_core.json"
    if not core_file.exists():
        print(f"{core_file} is missing: it lists the frozen core and the gold frames an older model trained on")
        return 2
    cj = json.loads(core_file.read_text(encoding="utf-8"))
    core, seen = set(cj["core"]), set(cj.get("seen", []))
    trained_on = set()
    for m in (root / "data" / "datasets").glob("*/manifest.csv"):
        with open(m, encoding="utf-8") as fh:
            next(fh, None)
            for line in fh:
                key = line.split(",")[0].strip()
                if key.startswith("tpi_"):
                    trained_on.add(key)
    groups = gate_stats.load_groups(gold.parent)
    stems = sorted(p.stem for p in gold.glob("*.txt") if p.stem not in groups["excluded"])
    key_of = lambda s: "tpi_" + re.search(r"tpi_(\d+)", s).group(1)
    unseen = [s for s in stems if s not in seen and key_of(s) not in trained_on]
    core_now = [s for s in unseen if s in core]

    ver2run = {}
    for rj in (root / "runs").glob("*/run.json"):
        r = json.loads(rj.read_text(encoding="utf-8"))
        if r.get("version"):
            ver2run[r["version"]] = r["id"]
    active = None
    ap_file = root / "models" / "active" / "production.json"
    if ap_file.exists():
        active = json.loads(ap_file.read_text(encoding="utf-8")).get("version")

    frames: Dict[str, Dict[str, gate_stats.FrameCounts]] = {}
    label_of: Dict[str, str] = {}
    # with kept_machine.json the evaluation ignores a model's own outlines object by object;
    # without it, a frame the model pre-labelled leaves its measure altogether (the older rule)
    object_level = gate_stats.kept_outlines_known(gold.parent)
    prelabel = {} if object_level else gate_stats.load_prelabel(gold.parent)
    own: Dict[str, int] = {}
    for d in sorted((root / "runs").glob("_measure_*/gold_predictions")):
        ver = d.parent.name[len("_measure_"):]
        label = ("beamline " + ver) if ver.startswith(BEAMLINE_PREFIX) else ver2run.get(ver, ver)
        measured = gate_stats.per_frame(d, gold, stems=unseen)
        # a frame the model pre-labelled holds its own outlines wherever the reviewer kept them:
        # it is "seen" by that model and leaves its measure (gate_stats.without_own_proposals)
        frames[label] = gate_stats.without_own_proposals(measured, ver, prelabel)
        own[label] = len(measured) - len(frames[label])
        label_of[ver] = label
    order = sorted(frames, key=lambda k: (not k.startswith("beamline"), k))
    reps = {k: report(frames[k]) for k in order}
    core_reps = {k: report({s: f for s, f in frames[k].items() if s in core}) for k in order}
    # frames whose light makes them unusable (saturated) are part of the measure; this is the same
    # table without them, to tell a model that got worse from a set that holds more bad light
    bad_light = {s for s, state in groups["light"].items() if state == "saturated"}
    usable_reps = {k: report({s: f for s, f in frames[k].items() if s in core and s not in bad_light}) for k in order}

    print(f"clean gold frames: {len(unseen)}" + (f" (core {len(core_now)})" if len(unseen) != len(core_now) else ""))
    if object_level:
        kept = sum(len(v) for v in __import__("src.eval.yolo_labels", fromlist=["x"]).read_kept_outlines(gold.parent).values())
        print(f"{kept} gold outlines are a model's proposal kept unchanged by the reviewer: not judged for any model (data/gold/kept_machine.json)")
    elif any(own.values()):
        print("frames a model pre-labelled itself, left out of its own measure: "
              + ", ".join(f"{k} {n}" for k, n in own.items() if n))
    elif not prelabel:
        print("!! data/gold/prelabel.json is missing (tools/gold_prelabel.py): a model is measured on frames it pre-labelled")
    table("", reps)
    short = [k for k in order if reps[k]["dataset"]["gold_images"] + own.get(k, 0) < len(unseen)]
    if short:
        print(f"!! measured on fewer frames than the current unseen set (re-measure them): {', '.join(short)}")
    if len(unseen) != len(core_now):
        table(f"\ncore, the {len(core_now)} frames frozen on {cj.get('frozen', '?')}:", core_reps)

    boot = {}
    base_label = label_of.get(active) if active else None
    if base_label and base_label in frames:
        print(f"\nagainst the active model ({base_label}), paired bootstrap over the frames both have, drawn by mount: difference [95 % interval], share of draws won")
        classes = {"crystal": NAMES.index("crystal")}
        for k in order:
            if k == base_label:
                continue
            b = gate_stats.paired_bootstrap(frames[k], frames[base_label], draws=args.draws, classes=classes, groups=groups["mount"])
            if not b.get("frames"):
                continue
            boot[k] = b
            f, m, c = b["f1"], b["mask_score"], b["per_class"]["crystal"]["f1"]
            print(f"{k:46} F1 {f['delta']:+.3f} [{f['lo']:+.3f}, {f['hi']:+.3f}] {f['win']:4.0%}   mask {m['delta']:+.3f} [{m['lo']:+.3f}, {m['hi']:+.3f}] {m['win']:4.0%}   crystal {c['delta']:+.3f} [{c['lo']:+.3f}, {c['hi']:+.3f}] {c['win']:4.0%}")

    # per frame and model: overall and crystal counts, so any two models can be compared head to head
    # with their noise interval (the monitor page runs the paired bootstrap in the browser)
    crystal = NAMES.index("crystal")

    def row(f):
        o = gate_stats._total([f], None); c = gate_stats._total([f], crystal)
        return [int(o[0]), int(o[1]), int(o[2]), round(o[3], 3), int(c[0]), int(c[1]), int(c[2]), round(c[3], 3)]
    per_frame = {"stems": unseen, "mount": [groups["mount"].get(st) or st for st in unseen],
                 "light": [groups["light"].get(st, "ok") for st in unseen], "fields": ["tp", "fp", "fn", "iou_sum", "crystal_tp", "crystal_fp", "crystal_fn", "crystal_iou_sum"],
                 "models": {k: [row(frames[k][st]) if st in frames[k] else None for st in unseen] for k in order}}
    (root / "runs" / "_unseen_per_frame.json").write_text(json.dumps(per_frame, separators=(",", ":")), encoding="utf-8")

    out = {"clean_frames": len(unseen), "reports": reps, "core_frames": len(core_now), "core_frozen": cj.get("frozen"),
           "core_reports": core_reps, "active": active, "bootstrap_vs_active": boot,
           "core_reports_usable_light": usable_reps, "saturated_frames": sorted(s for s in bad_light if s in core),
           "mounts": len({groups["mount"].get(s) or s for s in unseen}), "excluded": groups["excluded"]}
    (root / "runs" / "clean_gold_comparison.json").write_text(json.dumps(out, indent=2), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())