"""Check for near-duplicate images in training set"""
import sys, os
os.chdir("/home/user/goat")
import hashlib
from PIL import Image
import numpy as np
from collections import defaultdict

print("Checking image similarity/duplicates...", flush=True)

img_dir = 'Data/Detection_dataset/images/train'
files = sorted(os.listdir(img_dir))

def avg_hash(img, size=16):
    img = img.convert('L').resize((size, size), Image.LANCZOS)
    arr = np.array(img)
    avg = arr.mean()
    return ''.join(['1' if p > avg else '0' for p in arr.flatten()])

def hamming(h1, h2):
    return sum(c1 != c2 for c1, c2 in zip(h1, h2))

print("Computing hashes for {} images...".format(len(files)), flush=True)
hashes = {}
for i, f in enumerate(files):
    if i % 200 == 0:
        print("  {}/{}".format(i, len(files)), flush=True)
    img = Image.open(os.path.join(img_dir, f))
    hashes[f] = avg_hash(img)

groups = defaultdict(list)
for f in files:
    cam = f.split('_2025')[0]
    dn = 'Day' if 'Day' in f else 'Night'
    groups[(cam, dn)].append(f)

print("\nNear-duplicate analysis (hamming distance < 10):", flush=True)
total_dupes = 0
dupe_pairs = []

for group_key, group_files in sorted(groups.items()):
    group_dupes = 0
    for i in range(len(group_files)):
        for j in range(i+1, len(group_files)):
            d = hamming(hashes[group_files[i]], hashes[group_files[j]])
            if d < 10:
                group_dupes += 1
                if d < 5:
                    dupe_pairs.append((group_files[i], group_files[j], d))
    if group_dupes > 0:
        print("  {} {}: {} near-duplicates in {} images".format(
            group_key[0], group_key[1], group_dupes, len(group_files)), flush=True)
    total_dupes += group_dupes

print("\nTotal near-duplicate pairs (d<10): {}".format(total_dupes), flush=True)
print("\nVery similar pairs (d<5):", flush=True)
for f1, f2, d in dupe_pairs[:20]:
    print("  d={}: {} vs {}".format(d, f1, f2), flush=True)
if len(dupe_pairs) > 20:
    print("  ... and {} more".format(len(dupe_pairs)-20), flush=True)

print("\nGroup sizes:", flush=True)
for k in sorted(groups.keys()):
    print("  {}: {} images".format(k, len(groups[k])), flush=True)

print("\nDONE", flush=True)