#!/usr/bin/env python3
"""Find near-duplicate stills via perceptual hashing (dHash + aHash).

Pure CPU, fast: 3k images hash in a few minutes. Clusters images whose
dHash differs by <= --max-dist bits (default 10 of 64, a deliberately
paranoid threshold tuned for broadcast replays of the same shot, not
same scene different frame).

Outputs a cluster report CSV/JSON + console summary of groups >1.
With --move, relocates every non-keepsake of each group to --dupe-dir
(keepsake = smallest filename in group by default, or sharpest with
--keep-sharp). Nothing deleted; --move is reversible.

    python find_dupes.py -I <dir> [--move --dupe-dir <dir>/dupes]
"""

import argparse
import csv
import json
import os
import shutil
import sys
from collections import defaultdict


def parse_args():
    p = argparse.ArgumentParser(
        description=(__doc__ or "").split("\n")[0])
    p.add_argument("-I", "--input-dir", required=True)
    p.add_argument("--max-dist", type=int, default=6,
                   help="Max differing dHash bits (of 64) to cluster. "
                        "Broadcast replays of the identical shot differ "
                        "by only a few bits; verified at 10 that "
                        "genuinely-different frames can still be "
                        "within threshold (dHash floor), so default 6.")
    p.add_argument("--max-adist", type=int, default=None,
                   help="Optional second gate: max differing aHash bits. "
                        "Different hashes rarely agree twice; same frames "
                        "always do.")
    p.add_argument("--max-cdist", type=float, default=0.10,
                   help="Max L1 distance between normalized RGB color "
                        "histograms (0.128-bin/bh). Broadcast replays "
                        "are the SAME frames (near-zero color distance); "
                        "different shots differ >>0.10. This is the "
                        "signal that dHash lacks; set high to disable.")
    p.add_argument("--report", default=None, help="CSV report path.")
    p.add_argument("--move", action="store_true",
                   help="Move non-keepsakes to --dupe-dir.")
    p.add_argument("--dupe-dir", default=None,
                   help="Default: <input-dir>/duplicates")
    p.add_argument("--keep-sharp", action="store_true",
                   help="Keepsake = sharpest (Laplacian var) instead of "
                        "lowest filename.")
    return p.parse_args()


def hashes(img):
    import cv2
    import numpy as np
    g = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # color histogram signature (normalized L1 distance separates
    # same-shot replays (~0) from different shots (>>0.1) — the signal
    # dHash lacks: same gradient structure, different content)
    hist = np.zeros(3 * 64, dtype=np.float32)
    for c in range(3):
        h = cv2.calcHist([img], [c], None, [64], [0, 256]).flatten()
        hist[c*64:(c+1)*64] = h
    hist /= hist.sum()
    # dHash 8x8 gradient
    small = cv2.resize(g, (9, 8), interpolation=cv2.INTER_AREA)
    diff = small[:, 1:] > small[:, :-1]
    d = 0
    for bit, v in enumerate(diff.flatten()):
        if v:
            d |= 1 << bit
    # aHash 8x8 mean (secondary agreement check)
    small2 = cv2.resize(g, (8, 8), interpolation=cv2.INTER_AREA)
    m = small2.mean()
    a = 0
    for bit, v in enumerate((small2 > m).flatten()):
        if v:
            a |= 1 << bit
    # sharpness proxy for --keep-sharp
    sharp = cv2.Laplacian(g, cv2.CV_64F).var()
    return d, a, sharp, hist


def main():
    args = parse_args()
    import cv2
    import numpy as np

    exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
    files = [f for f in sorted(os.listdir(args.input_dir))
             if f.lower().endswith(exts)
             and os.path.isfile(os.path.join(args.input_dir, f))]

    if not files:
        print("No images found.")
        sys.exit(1)
    print(f"Hashing {len(files)} images ...")

    data = {}
    for i, f in enumerate(files):
        img = cv2.imread(os.path.join(args.input_dir, f))
        if img is None:
            continue
        d, a, sharp, hist = hashes(img)
        data[f] = (d, a, sharp, hist)
        if (i + 1) % 500 == 0:
            print(f"  {i+1}/{len(files)}", flush=True)

    names = list(data)
    # strict (complete-linkage) clustering — NO transitive chaining:
    # seed with the ungrouped image having most neighbours, take its
    # within-threshold neighbourhood, then keep only members pairwise
    # within threshold. Union-find chaining merges unrelated frames
    # after enough hops (verified failure: 90-frame 'group' of six
    # different scenes).
    n = len(names)
    print(f"Pairwise compare {n}x{n} (numpy hamming) ...")
    dmat = np.array([data[f][0] for f in names], dtype=np.uint64)
    # hamming distances via xor + popcount (work in int64; numpy has no
    # sane uint64 shift-add promotion)
    x = (dmat[:, None] ^ dmat[None, :]).astype(np.int64)
    ham = np.zeros(x.shape, dtype=np.int16)
    for b in range(64):
        ham += ((x >> b) & 1).astype(np.int16)

    # combined close matrix: dHash AND aHash AND color-histogram gates.
    # dHash alone has a false-positive floor (verified: distinct frames
    # at distance <=10); a second, differently-computed hash plus a
    # color signature is what separates 'same frame replayed' from
    # 'same show, different shot'.
    close = ham <= args.max_dist

    amat = np.array([data[f][1] for f in names], dtype=np.uint64)
    xa = (amat[:, None] ^ amat[None, :]).astype(np.int64)
    aham = np.zeros(xa.shape, dtype=np.int16)
    for b in range(64):
        aham += ((xa >> b) & 1).astype(np.int16)
    close &= aham <= (args.max_adist
                      if args.max_adist is not None else 8)

    hmat = np.stack([data[f][3] for f in names])  # N x 192 hist rows
    cdist = np.zeros((n, n), dtype=np.float32)
    for i in range(n):  # row-wise keeps it ~seconds and RAM-light
        cdist[i] = np.abs(hmat - hmat[i]).sum(axis=1)
    close &= cdist <= args.max_cdist
    np.fill_diagonal(close, False)

    # strict (complete-linkage) clustering — NO transitive chaining:
    # seed with the ungrouped image having most neighbours, take its
    # within-threshold neighbourhood, then keep only members pairwise
    # within threshold. Union-find chaining merges unrelated frames
    # after enough hops (verified failure: 90-frame 'group' of six
    # different scenes).
    order = sorted(range(n), key=lambda i: -close[i].sum())
    used = np.zeros(n, dtype=bool)
    clusters = []
    for seed in order:
        if used[seed] or not close[seed].any():
            continue
        cand = [seed] + [j for j in range(n)
                         if close[seed][j] and not used[j]]
        group = []
        for j in cand:  # keep only pairwise-close members
            if all(bool(close[j][k]) for k in group):
                group.append(j)
        if len(group) > 1:
            for j in group:
                used[j] = True
            clusters.append(group)
    clusters.sort(key=len, reverse=True)

    dupes_total = sum(len(g) - 1 for g in clusters)
    print(f"\n{len(clusters)} duplicate groups, {dupes_total} redundant "
          f"images (threshold {args.max_dist}/64 bits)")

    rows = []
    for gi, g in enumerate(clusters):
        members = [names[i] for i in g]
        if args.keep_sharp:
            members.sort(key=lambda f: -data[f][2])
        keepsake = members[0]
        print(f"\n  group {gi+1} ({len(members)} imgs) keep: {keepsake}")
        for m in members[:8]:
            print(f"    {'KEEP ' if m == keepsake else '     '}{m}")
        if len(members) > 8:
            print(f"    ... +{len(members)-8} more")
        rows.append({"group": gi + 1, "keepsake": keepsake,
                     "members": " | ".join(members)})

    if args.report:
        with open(args.report, "w", newline="") as fh:
            w = csv.DictWriter(fh, ["group", "keepsake", "members"])
            w.writeheader()
            w.writerows(rows)
        print(f"\nReport: {args.report}")

    if args.move:
        dest = args.dupe_dir or os.path.join(args.input_dir, "duplicates")
        os.makedirs(dest, exist_ok=True)
        moved = 0
        for r in rows:
            members = r["members"].split(" | ")
            for m in members:
                if m != r["keepsake"]:
                    src = os.path.join(args.input_dir, m)
                    if os.path.exists(src):
                        shutil.move(src, os.path.join(dest, m))
                        moved += 1
        print(f"Moved {moved} redundant images to {dest} "
              f"(report lists groups+keepsakes).")


if __name__ == "__main__":
    main()