#!/usr/bin/env python3
"""Standalone sanity check for exact-match *_exact_processed dirs.

Verifies, per dataset/split:
  - every STORED template candidate has a class-set EQUAL to the sample's classes (exact),
  - no sample matches itself (self-exclusion preserved),
  - exact-match coverage among foreground-bearing samples (how many got >=1 template),
and prints PASS/FAIL.

Usage:
  python sanity_check_exact.py \
      ACDC_1percent_exact_processed ACDC_5percent_exact_processed \
      ACDC_10percent_exact_processed ACDC_20percent_exact_processed \
      MMs_10percent_exact_processed
  (no args -> checks the default set above)
"""
import json, os, glob, sys
import numpy as np
from PIL import Image

RL = "/projects/p32685/RL_dataset"
DEFAULT = [
    "ACDC_1percent_exact_processed", "ACDC_5percent_exact_processed",
    "ACDC_10percent_exact_processed", "ACDC_20percent_exact_processed",
    "MMs_10percent_exact_processed",
]


def bank_sets(proc):
    out = {}
    for npz in glob.glob(os.path.join(proc, "template", "*.npz")):
        tid = os.path.splitext(os.path.basename(npz))[0]
        d = np.load(npz, allow_pickle=True)
        out[tid] = frozenset(int(c) for c in d.get("classes", []))
    return out


def mask_set(p):
    for c in [p, "/projects/p32685/" + p.lstrip("/")]:
        if os.path.exists(c):
            m = np.array(Image.open(c))
            return frozenset(int(v) for v in np.unique(m) if v != 0)
    return None


def check(name):
    proc = os.path.join(RL, name)
    if not os.path.isdir(proc):
        print(f"  {name}: MISSING"); return False
    bank = bank_sets(proc)
    all_pass = True
    for split in ["val_entries.json", "train_entries.json"]:
        path = os.path.join(proc, "backup_entries", split)
        if not os.path.exists(path):
            continue
        ent = json.load(open(path))
        n = with_t = ok = bad = selfm = 0
        for e in ent:
            sc = mask_set(e["label_mask"])
            if sc is None or len(sc) == 0:
                continue
            n += 1
            ss = os.path.splitext(os.path.basename(e["image"]))[0]
            tpl = e.get("template")
            tpl = json.loads(tpl) if isinstance(tpl, str) else tpl
            if not tpl:
                continue
            with_t += 1
            for t in tpl:
                if t["template_id"] == ss:
                    selfm += 1
                if bank.get(t["template_id"]) == sc:
                    ok += 1
                else:
                    bad += 1
        verdict = (bad == 0 and selfm == 0)
        all_pass &= verdict
        cov = 100.0 * with_t / max(n, 1)
        print(f"  {name} [{split:17s}] fg-N={n:5d}  with_template={with_t:5d} ({cov:5.1f}%)  "
              f"exact_ok={ok:5d}  non_exact={bad}  self_matches={selfm}  -> {'PASS' if verdict else 'FAIL'}")
    return all_pass


if __name__ == "__main__":
    names = sys.argv[1:] or DEFAULT
    print("=== EXACT-MATCH SANITY CHECK ===")
    overall = True
    for nm in names:
        overall &= check(nm)
    print("\nOVERALL:", "ALL PASS" if overall else "FAIL (see above)")
    sys.exit(0 if overall else 1)