#!/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=10,
                   help="Max differing dHash bits (of 64) to cluster. "
                        "0-6 near-identical, 8-12 replay-same-shot, "
                        ">14 starts eating costume/pose changes.")
    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>/dupes")
    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
    g = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 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


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 = hashes(img)
        data[f] = (d, a, sharp)
        if (i + 1) % 500 == 0:
            print(f"  {i+1}/{len(files)}", flush=True)

    names = list(data)
    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)

    # union-find clustering on distance threshold
    parent = list(range(n))

    def find(i):
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    def union(i, j):
        ri, rj = find(i), find(j)
        if ri != rj:
            parent[ri] = rj

    idx = np.argwhere(ham <= args.max_dist)
    for i, j in idx:
        if i < j:
            union(int(i), int(j))

    groups = defaultdict(list)
    for i in range(n):
        groups[find(i)].append(i)
    clusters = [g for g in groups.values() if len(g) > 1]
    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, "dupes")
        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()