#!/usr/bin/env python3
"""Which gold frames belong together, and which have unusable light -> ``data/gold/groups.json``.

    python tools/gold_groups.py                      # rebuild groups.json from data/gold
    python tools/gold_groups.py --sync-from <studio gold dir>   # take the studio's gold set first

* ``mount``: the studio writes a ``mount`` id on every gold frame of ``gold_manifest.json`` (same
  sample, or consecutive uploads of one centring). Frames of a mount are the same loop seen again,
  so the gate's bootstrap draws mounts, not frames (tools/gate_stats.py). Only mounts with more
  than one gold frame are listed.
* ``light``: ``saturated`` when more than half of the frame is blown out, ``dark`` when its mean
  grey level is under 30. Saturated frames stay in the measure (the beamline produces them); the
  comparison also reports the result without them.
* ``excluded``: the manifest's ``dropped`` frames (an exact duplicate of another gold frame, say).
  They are removed from data/gold and never measured.

On 2026-10-02: 203 frames, 7 mounts with 24 frames, 6 saturated, 29 dark, 1 dropped
(tpi_14688, the same image as tpi_14689 with a different label).
"""
from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(ROOT))
IMG_EXT = (".jpg", ".jpeg", ".png")

SATURATED_SHARE = 0.5      # share of pixels at or above 250
DARK_MEAN = 30.0           # mean grey level


def light_state(saturated_share: float, mean: float) -> str:
    if saturated_share > SATURATED_SHARE:
        return "saturated"
    if mean < DARK_MEAN:
        return "dark"
    return "ok"


def image_light(path: Path) -> str:
    import cv2
    g = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
    if g is None:
        return "ok"
    g = cv2.resize(g, (510, 510), interpolation=cv2.INTER_AREA)
    return light_state(float((g >= 250).mean()), float(g.mean()))


def build(gold: Path, *, with_light: bool = True) -> Dict:
    manifest = json.loads((gold / "gold_manifest.json").read_text(encoding="utf-8"))
    images = manifest.get("images") or {}
    by_mount: Dict[str, list] = {}
    for stem, info in images.items():
        if info.get("mount"):
            by_mount.setdefault(str(info["mount"]), []).append(stem)
    mount = {stem: m for m, stems in sorted(by_mount.items()) if len(stems) > 1 for stem in sorted(stems)}
    light: Dict[str, str] = {}
    if with_light:
        for stem in sorted(images):
            img = next((gold / "images" / f"{stem}{e}" for e in IMG_EXT if (gold / "images" / f"{stem}{e}").exists()), None)
            if img is not None:
                state = image_light(img)
                if state != "ok":
                    light[stem] = state
    return {
        "note": "mount: gold frames of one mount (drawn together by the gate's bootstrap); light: frames whose light is not "
                "usable; excluded: frames removed from the gold set. Built by tools/gold_groups.py",
        "updated": datetime.now().isoformat(timespec="seconds"),
        "frames": len(images), "mounts_with_several_frames": len({m for m in mount.values()}),
        "mount": mount, "light": light, "excluded": dict(manifest.get("dropped") or {}),
    }


def prune(gold: Path) -> int:
    """Remove the label and image of every frame the manifest no longer lists."""
    keep = set(json.loads((gold / "gold_manifest.json").read_text(encoding="utf-8")).get("images") or {})
    n = 0
    for sub in ("labels", "images"):
        for p in sorted((gold / sub).glob("*")) if (gold / sub).exists() else []:
            if p.stem.startswith("tpi_") and p.stem not in keep:
                p.unlink()
                n += 1
    return n


def main(argv: Optional[list] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=str(ROOT))
    ap.add_argument("--sync-from", default=None, help="studio gold directory (images/, labels/, gold_manifest.json) to import first")
    ap.add_argument("--no-light", action="store_true", help="skip reading the images")
    args = ap.parse_args(argv)
    root = Path(args.root).expanduser()
    gold = root / "data" / "gold"
    if args.sync_from:
        from tools.import_studio_export import import_gold
        info = import_gold(Path(args.sync_from).expanduser(), root=root)
        print(f"gold: {info['images']} images taken from {args.sync_from}")
    removed = prune(gold)
    if removed:
        print(f"gold: removed {removed} files of frames no longer in the manifest")
    g = build(gold, with_light=not args.no_light)
    (gold / "groups.json").write_text(json.dumps(g, indent=1), encoding="utf-8")
    states = {}
    for s in g["light"].values():
        states[s] = states.get(s, 0) + 1
    print(f"gold: {g['frames']} frames, {g['mounts_with_several_frames']} mounts holding {len(g['mount'])} of them, "
          f"light {states or 'all ok'}, excluded {len(g['excluded'])} -> {gold / 'groups.json'}")
    return 0


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