"""
audit_dedup.py — Duplicate / leakage audit for SpiceNet-Bench (Path A, step 2).

Pure-CPU near-duplicate audit using a 64-bit difference hash (dHash). No extra
deps (PIL + numpy only). Answers the three questions a Q1 reviewer asks of a
two-source benchmark whose headline is a cross-source accuracy gap:

  (L) TRAIN/TEST LEAKAGE  — is any test image a (near-)duplicate of a training
      image of the SAME source/class?  Leakage inflates the within-source ~99%.
  (X) CROSS-SOURCE LEAK   — within an overlap class, is any in-the-wild image a
      near-duplicate of a studio image?  A cross-source dup would confound the
      asymmetric -38 pp finding (the existential check).
  (D) WITHIN-SOURCE DUPS  — near-duplicate clusters inside one source (reported
      for the datasheet; not disqualifying on its own).

Outputs outputs/dedup_audit.json + a printed verdict.

    python audit_dedup.py
    python audit_dedup.py --base outputs/unified_benchmark.json --near 5
"""
import argparse
import json
from collections import defaultdict
from pathlib import Path

import numpy as np
from PIL import Image


def dhash(path, size=8):
    """64-bit difference hash as a Python int; None on failure."""
    try:
        img = Image.open(path).convert("L").resize((size + 1, size), Image.BILINEAR)
    except Exception:
        return None
    a = np.asarray(img, dtype=np.int16)
    diff = a[:, 1:] > a[:, :-1]                  # (size, size) bool
    packed = np.packbits(diff.flatten())         # 8 bytes
    return int.from_bytes(packed.tobytes(), "big")


def source_of(path):
    p = str(path).lower().replace("\\", "/")
    if "indian_spices" in p:
        return "indian"
    if "spice_spectrum" in p:
        return "spice_spectrum"
    return "unknown"


def popcount(arr_u64):
    """Vectorized set-bit count for a uint64 array."""
    return np.unpackbits(arr_u64.view(np.uint8).reshape(-1, 8), axis=1).sum(axis=1)


def near_pairs(ha, ia, hb, ib, thr, same=False):
    """(i, j, dist) with Hamming(a_i, b_j) <= thr. ha/hb uint64, ia/ib orig idx."""
    out = []
    if len(ha) == 0 or len(hb) == 0:
        return out
    for k in range(len(ha)):
        d = popcount(np.bitwise_xor(ha[k], hb))
        for jb in np.where(d <= thr)[0]:
            if same and ib[jb] <= ia[k]:
                continue                          # within one set: avoid self/double-count
            out.append((int(ia[k]), int(ib[jb]), int(d[jb])))
    return out


# dHash near-collision is necessary but NOT sufficient for "same image": on uniform
# studio backgrounds many DISTINCT photos share gross gradient structure. We confirm
# every candidate at the pixel level (32x32 grayscale, normalized RMSE) so the reported
# duplicate counts are defensible.
_PIXCACHE = {}


def small_gray(path, s=32):
    if path in _PIXCACHE:
        return _PIXCACHE[path]
    try:
        img = Image.open(path).convert("L").resize((s, s), Image.BILINEAR)
        a = np.asarray(img, dtype=np.float32) / 255.0
    except Exception:
        a = None
    _PIXCACHE[path] = a
    return a


def pixel_rmse(pa, pb):
    a, b = small_gray(pa), small_gray(pb)
    if a is None or b is None:
        return 1.0
    return float(np.sqrt(np.mean((a - b) ** 2)))


def verify(pairs, paths, rmse_thr):
    """Pixel-verify candidate pairs; return list of (i,j,hamming,rmse) with rmse<=thr."""
    keep = []
    for i, j, d in pairs:
        r = pixel_rmse(paths[i], paths[j])
        if r <= rmse_thr:
            keep.append((i, j, d, r))
    return keep


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="outputs/unified_benchmark.json")
    ap.add_argument("--near", type=int, default=5, help="Hamming <= near = near-duplicate")
    ap.add_argument("--out", default="outputs/dedup_audit.json")
    ap.add_argument("--max-examples", type=int, default=8)
    args = ap.parse_args()

    m = json.load(open(args.base))
    idx2name = {c["index"]: c["name"] for c in m["classes"]}
    paths, cls, src, split = [], [], [], []
    for sp in ("train", "val", "test"):
        for path, label in m["samples"][sp]:
            paths.append(path); cls.append(idx2name[int(label)])
            src.append(source_of(path)); split.append(sp)
    n = len(paths)
    cls = np.array(cls); src = np.array(src); split = np.array(split)

    print(f"hashing {n} images (dHash, CPU) ...")
    hashes = np.zeros(n, dtype=np.uint64)
    ok = np.ones(n, dtype=bool)
    for i, p in enumerate(paths):
        h = dhash(p)
        if h is None:
            ok[i] = False
        else:
            hashes[i] = np.uint64(h)
        if (i + 1) % 2000 == 0:
            print(f"  {i + 1}/{n}")
    print(f"  hashed ok={int(ok.sum())} fail={int((~ok).sum())}")

    idx = np.arange(n)
    classes = sorted(set(cls[ok]))
    src_by_cls = {c: sorted(set(src[ok & (cls == c)])) for c in classes}
    overlap = [c for c in classes if {"indian", "spice_spectrum"} <= set(src_by_cls[c])]
    print(f"classes={len(classes)} overlap(both-source)={overlap}")

    def grp(mask):
        ii = idx[ok & mask]
        return hashes[ii], ii

    # ── candidate pairs (loose dHash gate ≤ near), pixel-verified afterwards ─────
    x_cand, l_cand = [], []
    for c in overlap:                                  # (X) cross-source, overlap classes
        hi, ii = grp((cls == c) & (src == "indian"))
        hs, is_ = grp((cls == c) & (src == "spice_spectrum"))
        x_cand += near_pairs(hi, ii, hs, is_, args.near)
    for c in classes:                                  # (L) train/val -> test, same source+class
        for s in ("indian", "spice_spectrum"):
            ht, it = grp((cls == c) & (src == s) & (split == "test"))
            htr, itr = grp((cls == c) & (src == s) & (np.isin(split, ["train", "val"])))
            l_cand += near_pairs(ht, it, htr, itr, args.near)
    d_count = 0                                        # (D) within-source neighbours (informational)
    for c in classes:
        for s in ("indian", "spice_spectrum"):
            h, ii = grp((cls == c) & (src == s))
            d_count += len(near_pairs(h, ii, h, ii, args.near, same=True))

    buckets = defaultdict(list)                        # exact dHash collisions
    for i in idx[ok]:
        buckets[int(hashes[i])].append(int(i))
    exact_groups = [v for v in buckets.values() if len(v) > 1]
    exact_cross = sum(1 for g in exact_groups if len({src[i] for i in g}) > 1)

    print(f"candidates @ dHash<= {args.near}: cross-source={len(x_cand)} "
          f"leakage={len(l_cand)}; pixel-verifying {len(x_cand) + len(l_cand)} pairs ...")
    x_rows = [(i, j, d, pixel_rmse(paths[i], paths[j])) for i, j, d in x_cand]
    l_rows = [(i, j, d, pixel_rmse(paths[i], paths[j])) for i, j, d in l_cand]

    rmse_sweep = [0.02, 0.05, 0.10]                    # normalized-RMSE; 0.02 ~ re-encoded identical
    ham_sweep = [0, 1, 2, 3, args.near]
    DUP = 0.05                                         # pixel-confirmed duplicate threshold

    def le(rows, t):
        return sum(1 for *_, r in rows if r <= t)

    def ham_le(cand, t):
        return sum(1 for _, _, d in cand if d <= t)

    def ex(rows):
        rows = sorted([r for r in rows if r[3] <= DUP], key=lambda z: z[3])
        return [{"a": paths[i], "b": paths[j], "hamming": d, "rmse": round(r, 4),
                 "a_src": str(src[i]), "b_src": str(src[j]),
                 "a_split": str(split[i]), "b_split": str(split[j]), "class": str(cls[i])}
                for i, j, d, r in rows[: args.max_examples]]

    report = {
        "base": args.base, "n_images": n, "n_hash_fail": int((~ok).sum()),
        "n_classes": len(classes), "overlap_classes": overlap,
        "method": "dHash(64) candidate gate, then 32x32 grayscale normalized-RMSE pixel verify",
        "dup_rmse_threshold": DUP,
        "cross_source": {
            "candidates_dhash": len(x_cand),
            "candidates_by_hamming": {str(t): ham_le(x_cand, t) for t in ham_sweep},
            "verified_by_rmse": {str(t): le(x_rows, t) for t in rmse_sweep},
            "verified_dups": le(x_rows, DUP), "examples": ex(x_rows)},
        "train_test_leakage": {
            "candidates_dhash": len(l_cand),
            "candidates_by_hamming": {str(t): ham_le(l_cand, t) for t in ham_sweep},
            "verified_by_rmse": {str(t): le(l_rows, t) for t in rmse_sweep},
            "verified_dups": le(l_rows, DUP), "examples": ex(l_rows)},
        "within_source_visual_neighbours": d_count,
        "exact_dhash": {"groups": len(exact_groups), "cross_source_groups": exact_cross},
        "verdict": {"cross_source_clean": le(x_rows, DUP) == 0,
                    "no_train_test_leakage": le(l_rows, DUP) == 0},
    }
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    json.dump(report, open(args.out, "w"), indent=2)

    print("\n" + "=" * 72)
    print("DEDUP AUDIT VERDICT  (verified = pixel-confirmed, normalized RMSE <= %.2f)" % DUP)
    print("-" * 72)
    print("  CROSS-SOURCE (overlap classes) — would confound the -38pp finding if >0")
    print(f"    dHash candidates {len(x_cand)} | by Hamming {ham_sweep}: {[ham_le(x_cand,t) for t in ham_sweep]}")
    print(f"    pixel-verified by RMSE {rmse_sweep}: {[le(x_rows,t) for t in rmse_sweep]}")
    print(f"    => VERIFIED cross-source duplicates: {le(x_rows, DUP)}  "
          f"{'[CLEAN]' if le(x_rows,DUP)==0 else '[INVESTIGATE]'}")
    print("  TRAIN/VAL -> TEST LEAKAGE (same source+class) — would inflate ~99% if >0")
    print(f"    dHash candidates {len(l_cand)} | by Hamming {ham_sweep}: {[ham_le(l_cand,t) for t in ham_sweep]}")
    print(f"    pixel-verified by RMSE {rmse_sweep}: {[le(l_rows,t) for t in rmse_sweep]}")
    print(f"    => VERIFIED leakage duplicates: {le(l_rows, DUP)}  "
          f"{'[CLEAN]' if le(l_rows,DUP)==0 else '[INVESTIGATE]'}")
    print(f"  within-source visual neighbours (expected on studio data): {d_count}")
    print(f"  exact-dHash groups: {len(exact_groups)} (cross-source: {exact_cross})")
    print(f"  -> {args.out}")
    print("=" * 72)


if __name__ == "__main__":
    main()