#!/usr/bin/env -S uv run
# /// script
# dependencies = [
#   "torch",
#   "torchvision",
#   "timm",
#   "scikit-learn",
#   "pillow",
#   "imagehash",
# ]
# ///
"""Group similar large images using ResNet embeddings or perceptual hashing."""

import torch
import timm
import numpy as np
from PIL import Image
from pathlib import Path
from sklearn.cluster import KMeans, DBSCAN
from sklearn.preprocessing import StandardScaler
import imagehash
import argparse

# Load pretrained ResNet50 on M1
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
model = timm.create_model("resnet50", pretrained=True, num_classes=0)
model = model.to(device)
model.eval()


def get_embedding(img_path, size=224):
    """Extract embedding from image."""
    try:
        img = Image.open(img_path).convert("RGB")
        img.thumbnail((size, size), Image.Resampling.LANCZOS)

        canvas = Image.new("RGB", (size, size), (128, 128, 128))
        offset = ((size - img.width) // 2, (size - img.height) // 2)
        canvas.paste(img, offset)

        x = torch.tensor(np.array(canvas), dtype=torch.float32)
        x = x.permute(2, 0, 1) / 255.0
        x = (x - torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)) / torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

        with torch.no_grad():
            embedding = model(x.unsqueeze(0).to(device)).squeeze().cpu().numpy()
        return embedding
    except Exception as e:
        print(f"Error processing {img_path}: {e}")
        return None


def get_phash(img_path):
    """Extract perceptual hash from image."""
    try:
        img = Image.open(img_path).convert("RGB")
        return imagehash.phash(img)
    except Exception as e:
        print(f"Error processing {img_path}: {e}")
        return None


def hash_distance(h1, h2):
    """Hamming distance between two hashes."""
    return h1 - h2


def parse_args():
    parser = argparse.ArgumentParser(description="Group similar images using embeddings or perceptual hashing.")
    parser.add_argument("paths", nargs="+", type=Path, help="Image file(s) or directory")
    parser.add_argument("--cluster", action="store_true", help="Use ResNet embedding clustering")
    parser.add_argument("--perceptual-hash", action="store_true", help="Use perceptual hashing")
    parser.add_argument("--clusters", type=int, help="Number of clusters for embedding mode (auto if not specified)")
    parser.add_argument("--hash-threshold", type=int, default=5, help="Hamming distance threshold for perceptual hash")
    args = parser.parse_args()

    if not args.cluster and not args.perceptual_hash:
        parser.error("Either --cluster or --perceptual-hash must be specified")
    if args.cluster and args.perceptual_hash:
        parser.error("Cannot specify both --cluster and --perceptual-hash")

    return args


def group_by_hash(valid_files, hashes, threshold):
    """Group images by perceptual hash similarity."""
    labels = [-1] * len(valid_files)
    cluster_id = 0

    for i in range(len(valid_files)):
        if labels[i] != -1:
            continue
        labels[i] = cluster_id
        for j in range(i + 1, len(valid_files)):
            if labels[j] == -1 and hash_distance(hashes[i], hashes[j]) <= threshold:
                labels[j] = cluster_id
        cluster_id += 1

    return np.array(labels)


def main():
    args = parse_args()

    # Collect image files
    img_files = []
    for path in args.paths:
        if path.is_dir():
            img_files.extend(path.glob("*.[jJ][pP][gG]"))
            img_files.extend(path.glob("*.[pP][nN][gG]"))
        else:
            img_files.append(path)

    if not img_files:
        print("No images found.")
        return

    if args.perceptual_hash:
        # Perceptual hash mode
        hashes = []
        valid_files = []

        for i, img_path in enumerate(img_files):
            print(f"Processing {i + 1}/{len(img_files)}: {img_path.name}")
            h = get_phash(img_path)
            if h is not None:
                hashes.append(h)
                valid_files.append(img_path)

        labels = group_by_hash(valid_files, hashes, args.hash_threshold)
        n_clusters = len(np.unique(labels))

    else:
        # Embedding clustering mode
        embeddings = []
        valid_files = []

        for i, img_path in enumerate(img_files):
            print(f"Processing {i + 1}/{len(img_files)}: {img_path.name}")
            emb = get_embedding(img_path)
            if emb is not None:
                embeddings.append(emb)
                valid_files.append(img_path)

        embeddings = np.array(embeddings)
        scaler = StandardScaler()
        embeddings = scaler.fit_transform(embeddings)

        n_clusters = args.clusters or max(2, int(np.sqrt(len(embeddings) / 2)))
        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        labels = kmeans.fit_predict(embeddings)

    # Save clustered images
    for img_path, cluster_label in zip(valid_files, labels):
        new_name = f"group{cluster_label}_{img_path.name}"
        dest = Path.cwd() / new_name
        img_path.rename(dest) if img_path.parent == Path.cwd() else dest.write_bytes(img_path.read_bytes())

    print(f"Clustered {len(valid_files)} images into {n_clusters} clusters in {Path.cwd()}")


if __name__ == "__main__":
    main()