"""
make_dedup_manifests.py — leakage-free TEST splits (Path A, step 2).

For each input manifest, removes from the TEST split any image that is a
pixel-verified near-duplicate (dHash gate + 32x32 grayscale RMSE <= thr) of a
TRAIN or VAL image of the same class+source. TRAIN and VAL are left untouched,
so a checkpoint trained on the original split can be re-evaluated on the cleaned
test set WITHOUT retraining: its training distribution is unchanged, and the
remaining test images are provably absent (as near-duplicates) from train/val.

    python make_dedup_manifests.py --manifest outputs/manifest_overlap_indian.json
    python make_dedup_manifests.py --manifest outputs/manifest_overlap_ss.json
"""
import argparse
import json
from pathlib import Path

import numpy as np

from audit_dedup import dhash, near_pairs, pixel_rmse


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--near", type=int, default=5, help="dHash Hamming candidate gate")
    ap.add_argument("--rmse", type=float, default=0.05, help="pixel-verified duplicate threshold")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    m = json.load(open(args.manifest))
    items = []                                            # (path, label, split)
    for sp in ("train", "val", "test"):
        for path, label in m["samples"][sp]:
            items.append((path, int(label), sp))
    n = len(items)
    paths = [it[0] for it in items]
    labels = np.array([it[1] for it in items])
    splits = np.array([it[2] for it in items])

    print(f"{Path(args.manifest).name}: {n} images, hashing ...")
    H = 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:
            H[i] = np.uint64(h)

    idx = np.arange(n)
    n_test = int((splits == "test").sum())
    leaked = set()
    for c in sorted(set(labels.tolist())):
        te = idx[ok & (labels == c) & (splits == "test")]
        tr = idx[ok & (labels == c) & np.isin(splits, ["train", "val"])]
        for i, j, d in near_pairs(H[te], te, H[tr], tr, args.near):
            if pixel_rmse(paths[i], paths[j]) <= args.rmse:
                leaked.add(int(i))
    print(f"  leaked test images removed: {len(leaked)} / {n_test}")

    keep_test = [[paths[i], int(labels[i])] for i in idx
                 if splits[i] == "test" and i not in leaked]
    new = {k: v for k, v in m.items() if k != "samples"}
    new["samples"] = {"train": m["samples"]["train"], "val": m["samples"]["val"], "test": keep_test}
    new["dedup"] = {
        "method": f"dHash<={args.near} gate + 32x32 grayscale RMSE<={args.rmse} vs train/val (same class+source)",
        "removed_from_test": len(leaked), "test_before": n_test, "test_after": len(keep_test)}
    for cobj in new.get("classes", []):                  # refresh per-class test counts if present
        if "test" in cobj and "index" in cobj:
            ci = cobj["index"]
            cobj["test"] = sum(1 for _, l in keep_test if l == ci)
            cobj["total"] = cobj.get("train", 0) + cobj.get("val", 0) + cobj["test"]

    out = args.out or args.manifest.replace(".json", "_dedup.json")
    json.dump(new, open(out, "w"), indent=2)
    print(f"  -> {out}  (test {n_test} -> {len(keep_test)})")


if __name__ == "__main__":
    main()