#!/usr/bin/env python3
"""Does suppressing nested duplicate outlines help, and does it ever cost a real second crystal?

    python tools/dedup_study.py yolo26n-seg-v4-noflip-mask1-2026-10-02 [more versions] [--thresholds 0.5,0.7,0.9]

The models draw the same object twice (one outline inside the other) because the end-to-end
head has no NMS and IoU between a small and a large outline is low. Containment is not:
``intersection / smaller area``. The rule under study keeps, per class, the most confident
instance and drops any other instance of that class whose containment with a kept one is at
least the threshold. Two separate crystals have low containment and both stay.

For every model and threshold the script reports, on the gold frames no model trained on:
overall precision / recall / F1, the F1 of crystal, loop and face, how many instances the rule
dropped, and the crystal recall on the gold frames that hold two crystals or more (the frames
where a wrong rule would show). It also prints the containment between gold crystal pairs:
if human-drawn pairs never reach the threshold, the rule cannot lose one by construction.
Predictions are cached in runs/_dedup_<version>.json; results go to runs/dedup_study.json.
"""
from __future__ import annotations

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

import cv2
import numpy as np

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(ROOT))
from src.eval.metrics import ClassCount, aggregate, class_report, evaluate_image  # noqa: E402
from src.eval.yolo_labels import PolygonInstance, read_label_file  # noqa: E402

NAMES = ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
GRID = 512


def raster(points_norm: Sequence[Tuple[float, float]], size: int = GRID) -> np.ndarray:
    m = np.zeros((size, size), dtype=np.uint8)
    if len(points_norm) >= 3:
        pts = np.array([(x * size, y * size) for x, y in points_norm], dtype=np.int32).reshape(-1, 1, 2)
        cv2.fillPoly(m, [pts], 1)
    return m


def containment(a: np.ndarray, b: np.ndarray) -> float:
    """Intersection over the smaller area: 1 when one outline lies inside the other."""
    small = min(int(a.sum()), int(b.sum()))
    return float(np.logical_and(a, b).sum()) / small if small else 0.0


def suppress(insts: List[dict], thr: float, classes: set | None) -> Tuple[List[dict], int]:
    """Keep the most confident instance of a class; drop another of the same class contained in a
    kept one by ``thr`` or more. ``classes`` = class ids the rule applies to (None = all)."""
    kept: List[dict] = []
    masks: Dict[int, np.ndarray] = {}
    dropped = 0
    for inst in sorted(insts, key=lambda i: -i["conf"]):
        if classes is not None and inst["cls"] not in classes:
            kept.append(inst)
            continue
        m = raster(inst["points"])
        dup = False
        for k in kept:
            if k["cls"] != inst["cls"] or (classes is not None and k["cls"] not in classes):
                continue
            if containment(m, masks[id(k)]) >= thr:
                dup = True
                break
        if dup:
            dropped += 1
            continue
        masks[id(inst)] = m
        kept.append(inst)
    return kept, dropped


def predict_all(version: str, stems: List[str], img_dir: Path, conf: float = 0.25) -> Dict[str, List[dict]]:
    cache = ROOT / "runs" / f"_dedup_{version}.json"
    if cache.exists():
        return json.loads(cache.read_text())
    from ultralytics import YOLO
    model = YOLO(str(ROOT / "models" / "registry" / version / "weights.pt"))
    out: Dict[str, List[dict]] = {}
    for stem in stems:
        img = next((img_dir / f"{stem}{e}" for e in (".jpg", ".jpeg", ".png") if (img_dir / f"{stem}{e}").exists()), None)
        if img is None:
            continue
        res = model.predict(str(img), conf=conf, verbose=False)[0]
        insts = []
        if res.masks is not None:
            h, w = res.orig_shape
            for poly, cls, c in zip(res.masks.xy, res.boxes.cls.tolist(), res.boxes.conf.tolist()):
                if len(poly) >= 3:
                    insts.append({"cls": int(cls), "conf": float(c), "points": [(float(x) / w, float(y) / h) for x, y in poly]})
        out[stem] = insts
    cache.write_text(json.dumps(out))
    return out


def to_instances(insts: List[dict]) -> List[PolygonInstance]:
    return [PolygonInstance(i["cls"], [tuple(p) for p in i["points"]]) for i in insts]


def unseen_stems(gold_labels: Path) -> Tuple[List[str], set]:
    cj = json.loads((ROOT / "data" / "gold" / "unseen_core.json").read_text())
    seen = set(cj.get("seen", []))
    trained = 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.add(key)
    groups = ROOT / "data" / "gold" / "groups.json"
    excluded = set(json.loads(groups.read_text()).get("excluded", {})) if groups.exists() else set()
    key_of = lambda s: "tpi_" + re.search(r"tpi_(\d+)", s).group(1)
    stems = sorted(p.stem for p in gold_labels.glob("*.txt") if p.stem not in excluded)
    return [s for s in stems if s not in seen and key_of(s) not in trained], set(cj["core"])


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("versions", nargs="+")
    ap.add_argument("--thresholds", default="0.5,0.7,0.9")
    args = ap.parse_args(argv)
    thresholds = [float(t) for t in args.thresholds.split(",")]
    gold_labels, img_dir = ROOT / "data" / "gold" / "labels", ROOT / "data" / "gold" / "images"
    unseen, core = unseen_stems(gold_labels)
    gold = {s: read_label_file(gold_labels / f"{s}.txt") for s in unseen}
    crystal = NAMES.index("crystal")
    multi = [s for s in unseen if sum(1 for i in gold[s] if i.class_id == crystal) >= 2]

    # how close do human-drawn crystal pairs come to the rule?
    pair_c = []
    for s in multi:
        cs = [raster(i.points) for i in gold[s] if i.class_id == crystal]
        for a in range(len(cs)):
            for b in range(a + 1, len(cs)):
                pair_c.append(containment(cs[a], cs[b]))
    print(f"unseen gold frames {len(unseen)} (core {sum(1 for s in unseen if s in core)}), with two crystals or more: {len(multi)} "
          f"({sum(1 for s in multi for i in gold[s] if i.class_id == crystal)} crystals, {len(pair_c)} pairs)")
    if pair_c:
        pc = sorted(pair_c)
        print(f"containment between gold crystal pairs: median {pc[len(pc) // 2]:.2f}, max {pc[-1]:.2f}, "
              f"pairs at or above 0.5 / 0.7 / 0.9: {sum(c >= 0.5 for c in pc)} / {sum(c >= 0.7 for c in pc)} / {sum(c >= 0.9 for c in pc)}")

    results = {"unseen_frames": len(unseen), "multi_crystal_frames": len(multi), "gold_pair_containment": pair_c, "models": {}}
    variants = [("none", None, None)] + [(f"crystal>={t:g}", t, {crystal}) for t in thresholds] + [(f"all>={t:g}", t, None) for t in thresholds]
    for version in args.versions:
        preds = predict_all(version, unseen, img_dir)
        rows = {}
        print(f"\n== {version}")
        print(f"{'rule':14} {'P':>6} {'R':>6} {'F1':>6}  {'crystal':>7} {'loop':>6} {'face':>6} {'pin':>6}  {'dropped':>7}  {'crystal R on multi-crystal frames':>34}")
        for name, thr, classes in variants:
            per_image, dropped, multi_counts = [], 0, []
            for s in unseen:
                insts = preds.get(s, [])
                if thr is not None:
                    insts, d = suppress(insts, thr, classes)
                    dropped += d
                ev = evaluate_image(gold[s], to_instances(insts))
                per_image.append(ev)
                if s in multi:
                    multi_counts.append(ev.get(crystal, ClassCount()))
            total = aggregate(per_image)
            overall = ClassCount()
            for cc in total.values():
                overall.tp += cc.tp; overall.fp += cc.fp; overall.fn += cc.fn
            o = class_report(overall)
            per = {NAMES[c]: class_report(cc) for c, cc in total.items() if 0 <= c < len(NAMES)}
            mt = sum(c.tp for c in multi_counts); mf = sum(c.fn for c in multi_counts)
            rows[name] = {"precision": o["precision"], "recall": o["recall"], "f1": o["f1"], "per_class_f1": {k: v["f1"] for k, v in per.items()},
                          "fp": overall.fp, "dropped": dropped, "multi_crystal": {"tp": mt, "fn": mf}}
            f = lambda k: per.get(k, {}).get("f1", 0.0)
            print(f"{name:14} {o['precision']:6.3f} {o['recall']:6.3f} {o['f1']:6.3f}  {f('crystal'):7.3f} {f('loop_all'):6.3f} {f('loop_face'):6.3f} {f('pin'):6.3f}  {dropped:7d}  {mt:3d} / {mt + mf:3d}")
        results["models"][version] = rows
    (ROOT / "runs" / "dedup_study.json").write_text(json.dumps(results, indent=1))
    return 0


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