"""Detect near-duplicate images shared across train/valid/test splits.

The PPE dataset is built partly from video frames, and consecutive frames that
land in different splits leak training data into evaluation — inflating every
reported metric. This script perceptual-hashes every image and reports
cross-split pairs whose hashes are within a Hamming-distance threshold.

Needs a local dataset copy (scripts/download_data.py) and the `audit` group:

    uv run --group audit python scripts/check_split_leakage.py
    uv run --group audit python scripts/check_split_leakage.py --threshold 8 --csv leakage.csv

Threshold guide (64-bit pHash): 0 = pixel-identical or re-encoded copies,
<=4 = near-certain duplicates (crops / adjacent video frames), <=8 = likely
same scene. Start at the default 4; raise to 8 to gauge how bad it could be.
"""

from __future__ import annotations

import argparse
import csv
from itertools import combinations
from pathlib import Path

import imagehash
import numpy as np
from PIL import Image

SPLITS = ["train", "valid", "test"]
IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def hash_split(split_dir: Path) -> tuple[list[str], np.ndarray]:
    """pHash every image under split_dir/images -> (filenames, (n, 64) bool array)."""
    files = sorted(f for f in (split_dir / "images").iterdir() if f.suffix.lower() in IMG_EXTS)
    hashes = np.empty((len(files), 64), dtype=bool)
    for i, f in enumerate(files):
        with Image.open(f) as im:
            hashes[i] = imagehash.phash(im).hash.reshape(-1)
    return [f.name for f in files], hashes


def cross_split_pairs(
    ha: np.ndarray, hb: np.ndarray, threshold: int, chunk: int = 512
) -> list[tuple[int, int, int]]:
    """Index pairs (i, j, distance) with Hamming distance <= threshold, chunked to bound memory."""
    pairs = []
    for start in range(0, len(ha), chunk):
        block = ha[start : start + chunk]
        # (chunk, n_b) Hamming distances via broadcast XOR-count.
        dists = (block[:, None, :] != hb[None, :, :]).sum(axis=2)
        for bi, j in zip(*np.nonzero(dists <= threshold), strict=True):
            pairs.append((start + int(bi), int(j), int(dists[bi, j])))
    return pairs


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--data-dir", default="work", help="directory containing train/ valid/ test/")
    p.add_argument("--threshold", type=int, default=4, help="max Hamming distance to flag (of 64 bits)")
    p.add_argument("--csv", default=None, help="optional path to write the flagged pairs as CSV")
    args = p.parse_args()

    data_dir = Path(args.data_dir)
    splits: dict[str, tuple[list[str], np.ndarray]] = {}
    for split in SPLITS:
        if not (data_dir / split / "images").is_dir():
            print(f"skip {split}: {data_dir / split / 'images'} not found")
            continue
        names, hashes = hash_split(data_dir / split)
        splits[split] = (names, hashes)
        print(f"{split}: hashed {len(names)} images")

    if len(splits) < 2:
        raise SystemExit("Need at least two splits present — fetch the dataset with scripts/download_data.py")

    rows: list[tuple[str, str, str, str, int]] = []
    for sa, sb in combinations(splits, 2):
        names_a, ha = splits[sa]
        names_b, hb = splits[sb]
        pairs = cross_split_pairs(ha, hb, args.threshold)
        print(f"\n{sa} × {sb}: {len(pairs)} pairs within Hamming distance {args.threshold}")
        for i, j, d in sorted(pairs, key=lambda t: t[2])[:20]:
            print(f"  d={d:2d}  {sa}/{names_a[i]}  ~  {sb}/{names_b[j]}")
        if len(pairs) > 20:
            print(f"  ... {len(pairs) - 20} more (use --csv for the full list)")
        rows += [(sa, names_a[i], sb, names_b[j], d) for i, j, d in pairs]

    # The verdict that matters: what fraction of each eval split is contaminated by train?
    print("\n--- Summary ---")
    for eval_split in ("valid", "test"):
        if eval_split not in splits or "train" not in splits:
            continue
        contaminated = {img_b for sa, _img_a, sb, img_b, _d in rows if sa == "train" and sb == eval_split}
        n = len(splits[eval_split][0])
        pct = 100 * len(contaminated) / n if n else 0.0
        print(f"{eval_split}: {len(contaminated)}/{n} images ({pct:.1f}%) have a near-duplicate in train")
    if rows:
        print("Metrics computed on contaminated splits overstate real-world performance.")
        print("Fix: move each duplicate group entirely into one split, then re-run eval.")
    else:
        print("No cross-split near-duplicates at this threshold. 🎉")

    if args.csv:
        with open(args.csv, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["split_a", "image_a", "split_b", "image_b", "hamming"])
            w.writerows(rows)
        print(f"\nWrote {len(rows)} pairs to {args.csv}")


if __name__ == "__main__":
    main()