import gradio as gr
import os
import cv2
import shutil
import time
from PIL import Image
from ultralytics import YOLO
import imagehash
import threading
import queue
from concurrent.futures import ProcessPoolExecutor, as_completed
import random
import string
import argparse
import logging
from tqdm import tqdm
import json
from pathlib import Path
import numpy as np

import torch  # Add this line

# Import SAM2 functions and constants
from SAM_Segmenter import (
    build_sam2, SAM2ImagePredictor, load_model,
    SAM2_CHECKPOINT, SAM2_MODEL_CONFIG,
    GROUNDING_DINO_CONFIG, GROUNDING_DINO_CHECKPOINT,
    process_image as sam2_process_image
)

Image.MAX_IMAGE_PIXELS = None
logging.basicConfig(level=logging.INFO)

os.environ['YOLO_AUTOINSTALL'] = '0'
os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'

img_formats = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.tif', '.webp', '.heic', '.raw',
'.cr2', '.nef', '.arw', '.svg', '.ico', '.psd', '.ai', '.eps', '.pdf', '.indd',
'.jfif', '.pjpeg', '.pjp', '.hdr', '.exr', '.dib', '.svgz', '.avif', '.dng',
'.orf', '.rw2', '.sr2', '.3fr', '.srf', '.psb', '.emf', '.wmf', '.cdr',
'.cgm', '.pgm', '.ppm', '.pbm', '.pnm', '.sgi', '.ras', '.xbm', '.xpm')

YOLO_CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane',
    'bus', 'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
    'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
    'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
    'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle',
    'wine glass', 'cup', 'fork', 'knife', 'spoon',
    'bowl', 'banana', 'apple', 'sandwich', 'orange',
    'broccoli', 'carrot', 'hot dog', 'pizza', 'donut',
    'cake', 'chair', 'couch', 'potted plant', 'bed',
    'dining table', 'toilet', 'tv', 'laptop', 'mouse',
    'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
    'toaster', 'sink', 'refrigerator', 'book', 'clock',
    'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

def resize_bbox_to_dimensions(bbox, target_width, target_height, img_width, img_height):
    x1, y1, x2, y2 = bbox
    current_width = x2 - x1
    current_height = y2 - y1
    desired_aspect_ratio = target_width / target_height
    current_aspect_ratio = current_width / current_height

    if current_aspect_ratio < desired_aspect_ratio:
        new_width = desired_aspect_ratio * current_height
        x1 -= (new_width - current_width) / 2
        x2 += (new_width - current_width) / 2
    elif current_aspect_ratio > desired_aspect_ratio:
        new_height = current_width / desired_aspect_ratio
        y1 -= (new_height - current_height) / 2
        y2 += (new_height - current_height) / 2

    x1 = max(x1, 0)
    y1 = max(y1, 0)
    x2 = min(x2, img_width)
    y2 = min(y2, img_height)

    return [int(x1), int(y1), int(x2), int(y2)]

def crop_images(input_folder, output_folder, aspect_ratios, yolo_folder, save_yolo, num_threads, overwrite, selected_class, save_as_png, sam2_prompt):
    aspect_ratios = [tuple(map(int, ar.strip().split('x'))) for ar in aspect_ratios.split(',')]
    os.makedirs(output_folder, exist_ok=True)
    if save_yolo:
        os.makedirs(yolo_folder, exist_ok=True)

    image_files = [f for f in os.listdir(input_folder) if f.lower().endswith(img_formats)]
    total_files = len(image_files)

    if sam2_prompt:
        # Use SAM2 with single thread
        sam2_model = build_sam2(SAM2_MODEL_CONFIG, SAM2_CHECKPOINT, device="cuda" if torch.cuda.is_available() else "cpu")
        sam2_predictor = SAM2ImagePredictor(sam2_model)
        grounding_model = load_model(
            model_config_path=GROUNDING_DINO_CONFIG,
            model_checkpoint_path=GROUNDING_DINO_CHECKPOINT,
            device="cuda" if torch.cuda.is_available() else "cpu"
        )
        
        results = []
        for filename in tqdm(image_files, total=total_files, desc="Processing images with SAM2"):
            result = process_image((filename, input_folder, output_folder, aspect_ratios, yolo_folder, save_yolo, overwrite, selected_class, save_as_png, sam2_prompt, None, sam2_predictor, grounding_model))
            results.append(result)
    else:
        # Use YOLO with multi-threading
        model = YOLO("models/yolo11x.pt")
        
        with ProcessPoolExecutor(max_workers=num_threads) as executor:
            args_list = [(filename, input_folder, output_folder, aspect_ratios, yolo_folder, save_yolo, overwrite, selected_class, save_as_png, sam2_prompt, model, None, None) for filename in image_files]
            results = list(tqdm(executor.map(process_image, args_list), total=total_files, desc="Processing images with YOLO"))

    processed_count = sum(r[1] for r in results)
    skipped_count = sum(r[2] for r in results)

    return f"Completed! Processed: {processed_count}, Skipped: {skipped_count}, Total: {total_files}"

from PIL import Image, ExifTags

def add_padding_to_bbox(bbox, img_width, img_height, padding_percent=0.1):
    x1, y1, x2, y2 = bbox
    width = x2 - x1
    height = y2 - y1
    
    padding_x = width * padding_percent
    padding_y = height * padding_percent
    
    x1 = max(0, x1 - padding_x)
    y1 = max(0, y1 - padding_y)
    x2 = min(img_width, x2 + padding_x)
    y2 = min(img_height, y2 + padding_y)
    
    return [int(x1), int(y1), int(x2), int(y2)]

def process_image(args):
    filename, input_folder, output_folder, aspect_ratios, yolo_folder, save_yolo, overwrite, selected_class, save_as_png, sam2_prompt, model, sam2_predictor, grounding_model = args
    try:
        img_path = os.path.join(input_folder, filename)
        img = Image.open(img_path)

        # Preserve EXIF data
        exif = img._getexif()
        if exif:
            exif = {ExifTags.TAGS[k]: v for k, v in exif.items() if k in ExifTags.TAGS}
            orientation = exif.get('Orientation', 1)
            
            # Rotate the image according to EXIF orientation
            if orientation == 3:
                img = img.rotate(180, expand=True)
            elif orientation == 6:
                img = img.rotate(270, expand=True)
            elif orientation == 8:
                img = img.rotate(90, expand=True)

        img = img.convert("RGB")
        img_array = np.array(img)

        if sam2_prompt:
            # Use SAM2 for segmentation
            segmented_image, mask_image, mask_coordinates, _ = sam2_process_image(
                img_array, sam2_prompt, 0.35, 0.25, True,
                sam2_predictor, grounding_model, output_folder, yolo_folder,
                filename, save_as_png, overwrite, img_path
            )
            if mask_coordinates:
                # Find the extreme points to create a bounding box
                x_coords = [p[0] for p in mask_coordinates]
                y_coords = [p[1] for p in mask_coordinates]
                x1, x2 = min(x_coords), max(x_coords)
                y1, y2 = min(y_coords), max(y_coords)
                
                # Add padding to ensure we don't crop too tightly
                padding_x = (x2 - x1) * 0.1  # 10% padding
                padding_y = (y2 - y1) * 0.1  # 10% padding
                
                x1 = max(0, x1 - padding_x)
                y1 = max(0, y1 - padding_y)
                x2 = min(img.width, x2 + padding_x)
                y2 = min(img.height, y2 + padding_y)
                
                bbox = [int(x1), int(y1), int(x2), int(y2)]
            else:
                return f"No segmentation found for {filename}", 0, 1
        else:
            # Use YOLO for detection
            results = model(img_array)
            detections = results[0].boxes
            class_detections = [box for box in detections if model.names[int(box.cls)] == selected_class]
            if class_detections:
                bbox = class_detections[0].xyxy[0].tolist()
                # Add padding to YOLO bbox to ensure the entire object is included
                bbox = add_padding_to_bbox(bbox, img.width, img.height)
            else:
                return f"No {selected_class} detected in {filename}", 0, 1

        processed = 0
        skipped = 0
        for target_width, target_height in aspect_ratios:
            new_bbox = resize_bbox_to_dimensions(bbox, target_width, target_height, img.width, img.height)
            cropped_img = img.crop(new_bbox)

            aspect_folder = os.path.join(output_folder, f"{target_width}x{target_height}")
            os.makedirs(aspect_folder, exist_ok=True)

            save_path = os.path.join(aspect_folder, os.path.splitext(filename)[0] + ('.png' if save_as_png else '.jpg'))

            if not overwrite and os.path.exists(save_path):
                skipped += 1
                continue

            # Save the cropped image with original EXIF data
            if 'exif' in img.info:
                cropped_img.save(save_path, format='PNG' if save_as_png else 'JPEG', quality=100, exif=img.info['exif'])
            else:
                cropped_img.save(save_path, format='PNG' if save_as_png else 'JPEG', quality=100)
            processed += 1

        if save_yolo and yolo_folder:
            yolo_save_path = os.path.join(yolo_folder, os.path.splitext(filename)[0] + ('.png' if save_as_png else '.jpg'))
            if overwrite or not os.path.exists(yolo_save_path):
                if sam2_prompt:
                    Image.fromarray(segmented_image).save(yolo_save_path, format='PNG' if save_as_png else 'JPEG', quality=100)
                else:
                    Image.fromarray(results[0].plot()).save(yolo_save_path, format='PNG' if save_as_png else 'JPEG', quality=100)
                processed += 1
            else:
                skipped += 1

        return f"Processed {filename}", processed, skipped
    except Exception as e:
        return f"Error processing {filename}: {str(e)}", 0, 1
    
def resize_image(args):
    input_path, resolutions, output_root, save_as_png, face_model_path, overwrite, no_crop = args
    try:
        image = Image.open(input_path)
        processed = 0
        skipped = 0
        face_detected = False
        face_locations = None

        for resolution in resolutions:
            subfolder_name = f"{resolution[0]}x{resolution[1]}"
            output_folder = os.path.join(output_root, subfolder_name)
            os.makedirs(output_folder, exist_ok=True)

            output_filename = os.path.basename(input_path)
            if save_as_png:
                output_filename = os.path.splitext(output_filename)[0] + '.png'
            else:
                output_filename = os.path.splitext(output_filename)[0] + '.jpg'
            output_path = os.path.join(output_folder, output_filename)

            if not overwrite and os.path.exists(output_path):
                logging.info(f"Skipping {output_filename} for {resolution[0]}x{resolution[1]} as it already exists.")
                skipped += 1
                continue

            if no_crop:
                new_image = Image.new('RGB', resolution, (255, 255, 255))
                aspect_ratio = min(resolution[0] / image.width, resolution[1] / image.height)
                new_size = (int(image.width * aspect_ratio), int(image.height * aspect_ratio))
                resized_image = image.resize(new_size, Image.LANCZOS)
                position = ((resolution[0] - new_size[0]) // 2, (resolution[1] - new_size[1]) // 2)
                new_image.paste(resized_image, position)
                resized_image = new_image
            else:
                if not face_detected:
                    face_model = YOLO(face_model_path)
                    results = face_model(input_path)
                    detections = results[0].boxes

                    if len(detections) > 0:
                        face_detected = True
                        face = detections[0]
                        x1, y1, x2, y2 = face.xyxy[0].tolist()
                        face_locations = [(int(y1), int(x2), int(y2), int(x1))]

                desired_aspect_ratio = resolution[0] / resolution[1]
                image_aspect_ratio = image.width / image.height

                if image_aspect_ratio > desired_aspect_ratio:
                    new_width = int(image.height * desired_aspect_ratio)
                    new_height = image.height
                else:
                    new_width = image.width
                    new_height = int(image.width / desired_aspect_ratio)

                left = (image.width - new_width) / 2
                top = (image.height - new_height) / 2
                right = (image.width + new_width) / 2
                bottom = (image.height + new_height) / 2

                if face_detected and face_locations:
                    face_top, face_right, face_bottom, face_left = face_locations[0]
                    face_center_x = (face_left + face_right) // 2
                    face_center_y = (face_top + face_bottom) // 2

                    left = min(max(0, face_center_x - new_width // 2), image.width - new_width)
                    top = min(max(0, face_center_y - new_height // 2), image.height - new_height)
                    right = left + new_width
                    bottom = top + new_height

                resized_image = image.crop((left, top, right, bottom))
                resized_image = resized_image.resize(resolution, Image.LANCZOS)

            if save_as_png:
                resized_image.save(output_path, format='PNG')
            else:
                resized_image.save(output_path, format='JPEG', quality=100)

            processed += 1

        return processed, skipped
    except Exception as e:
        logging.error(f"Error processing {input_path}: {e}")
        return 0, 0

def resize_images(input_folder, output_folder, resolutions, save_as_png, num_threads=1, overwrite=False, no_crop=False):
    resolutions = [tuple(map(int, res.strip().split('x'))) for res in resolutions.split(',')]
    face_model_path = "models/face_yolov9c.pt"

    image_paths = []
    for resolution in resolutions:
        resolution_folder = os.path.join(input_folder, f"{resolution[0]}x{resolution[1]}")
        if os.path.exists(resolution_folder):
            resolution_images = [
                os.path.join(resolution_folder, fname)
                for fname in os.listdir(resolution_folder)
                if fname.lower().endswith(img_formats) and os.path.isfile(os.path.join(resolution_folder, fname))
            ]
            image_paths.extend(resolution_images)
            logging.info(f"Found {len(resolution_images)} images in {resolution_folder}")
        else:
            error_msg = f"Folder not found for resolution {resolution[0]}x{resolution[1]}. Skipping this resolution."
            logging.error(error_msg)
            yield error_msg

    total_images = len(image_paths)
    logging.info(f"Total images to process: {total_images}")

    if total_images == 0:
        yield "No images found to process in the specified resolution folders."
        return

    num_threads = min(num_threads, total_images)
    logging.info(f"Using {num_threads} threads for processing")

    processed_count = 0
    skipped_count = 0
    lock = threading.Lock()
    start_time = time.time()

    def update_count(processed, skipped):
        nonlocal processed_count, skipped_count
        with lock:
            processed_count += processed
            skipped_count += skipped
            return processed_count + skipped_count

    progress_bar = tqdm(total=total_images, unit='image', desc="Resizing images")

    with ProcessPoolExecutor(max_workers=num_threads) as executor:
        future_to_image = {executor.submit(resize_image, (path, resolutions, output_folder, save_as_png, face_model_path, overwrite, no_crop)): path
                           for path in image_paths}

        for future in as_completed(future_to_image):
            try:
                processed, skipped = future.result()
                progress_bar.update(processed + skipped)
                count = update_count(processed, skipped)
                elapsed_time = time.time() - start_time
                images_per_second = count / elapsed_time if elapsed_time > 0 else 0
                eta_seconds = (total_images - count) / images_per_second if images_per_second > 0 else 0
                eta_str = time.strftime("%H:%M:%S", time.gmtime(eta_seconds))

                progress_msg = f"Processed: {processed_count}, Skipped: {skipped_count}, Total: {count}/{total_images} images. ETA: {eta_str}"
                print(progress_msg)  # This will show in the command line
                yield progress_msg  # This will update the Gradio interface

                if count >= total_images:
                    break
            except Exception as exc:
                logging.error(f"Generated an exception: {exc}")

    progress_bar.close()
    total_time = time.time() - start_time
    yield f"Processing complete! Processed: {processed_count}, Skipped: {skipped_count}, Total: {processed_count + skipped_count}/{total_images} images in {time.strftime('%H:%M:%S', time.gmtime(total_time))}"

def move_low_res_files(folder_a, folder_b, folder_c, folder_d, min_width, min_height):
    os.makedirs(folder_c, exist_ok=True)
    os.makedirs(folder_d, exist_ok=True)

    image_files = [f for f in os.listdir(folder_a) if f.endswith(img_formats)]
    total_files = len(image_files)
    files_processed = 0
    files_moved = 0

    for filename in image_files:
        image_path_a = os.path.join(folder_a, filename)
        image_path_b = os.path.join(folder_b, filename)

        if os.path.exists(image_path_b):
            with Image.open(image_path_a) as img:
                width, height = img.size

            if width < min_width or height < min_height:
                shutil.move(image_path_b, folder_c)
                shutil.move(image_path_a, folder_d)
                files_moved += 1

        files_processed += 1
        yield f"Processed: {files_processed}/{total_files}, Moved: {files_moved}"

    yield f"Processing complete! Processed: {files_processed}/{total_files}, Moved: {files_moved}"

def rename_files(folder_path, start_number):
    def random_string(length=10):
        return ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(length))

    def assign_random_names(path):
        for filename in os.listdir(path):
            file_path = os.path.join(path, filename)
            if os.path.isfile(file_path):
                random_name = f"{random_string()}.{filename.split('.')[-1]}"
                random_file_path = os.path.join(path, random_name)
                os.rename(file_path, random_file_path)
                yield f"Assigned random name: {random_name}"

    def rename_files_sequentially(path, start):
        i = start
        for filename in os.listdir(path):
            file_path = os.path.join(path, filename)
            if os.path.isfile(file_path):
                new_filename = f"img{i}.{filename.split('.')[-1]}"
                new_file_path = os.path.join(path, new_filename)
                os.rename(file_path, new_file_path)
                i += 1
                yield f"Renamed to: {new_filename}"

    yield from assign_random_names(folder_path)
    yield from rename_files_sequentially(folder_path, start_number)
    yield "Renaming complete!"

def process_image_face(args):
    input_path, output_root, padding, save_as_png, face_model = args
    try:
        # Convert image to PNG using PIL
        with Image.open(input_path) as img:
            png_path = os.path.splitext(input_path)[0] + '_temp.png'
            img.save(png_path, 'PNG')

        # Read the PNG with cv2
        image = cv2.imread(png_path)

        if image is None:
            raise ValueError(f"Could not read image {input_path} even after conversion to PNG")

        results = face_model(png_path)
        detections = results[0].boxes

        if len(detections) == 0:
            os.remove(png_path)  # Clean up temporary PNG file
            return f"No face detected in {input_path}", 0, 1  # Skipped

        face = detections[0]
        x1, y1, x2, y2 = map(int, face.xyxy[0].tolist())

        # Apply padding
        x1 = max(0, x1 - padding)
        y1 = max(0, y1 - padding)
        x2 = min(image.shape[1], x2 + padding)
        y2 = min(image.shape[0], y2 + padding)

        face_image = image[y1:y2, x1:x2]

        os.makedirs(output_root, exist_ok=True)
        output_filename = os.path.splitext(os.path.basename(input_path))[0] + ('.png' if save_as_png else '.jpg')
        output_path = os.path.join(output_root, output_filename)

        if save_as_png:
            cv2.imwrite(output_path, face_image)
        else:
            cv2.imwrite(output_path, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), 100])

        os.remove(png_path)  # Clean up temporary PNG file
        return f"Processed: {input_path}", 1, 0  # Processed
    except Exception as e:
        if 'png_path' in locals():
            os.remove(png_path)  # Clean up temporary PNG file if it exists
        return f"Error processing {input_path}: {e}", 0, 0  # Neither processed nor skipped

def extract_faces(input_folder, output_folder, padding, save_as_png, num_threads):
    face_model = YOLO("models/face_yolov9c.pt")

    image_paths = [os.path.join(input_folder, fname) for fname in os.listdir(input_folder) if fname.lower().endswith(img_formats)]
    total_images = len(image_paths)
    processed_count = 0
    skipped_count = 0
    start_time = time.time()
    lock = threading.Lock()

    def update_progress(processed, skipped):
        nonlocal processed_count, skipped_count
        with lock:
            processed_count += processed
            skipped_count += skipped
            return processed_count + skipped_count

    progress_bar = tqdm(total=total_images, unit='image', desc="Extracting faces")

    with ProcessPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(process_image_face, (path, output_folder, padding, save_as_png, face_model)) for path in image_paths]
        for future in as_completed(futures):
            result, processed, skipped = future.result()
            progress_bar.update(processed + skipped)
            total_processed = update_progress(processed, skipped)
            elapsed_time = time.time() - start_time
            images_per_second = total_processed / elapsed_time if elapsed_time > 0 else 0
            images_left = total_images - total_processed
            eta_seconds = images_left / images_per_second if images_per_second > 0 else 0
            eta_str = time.strftime("%H:%M:%S", time.gmtime(eta_seconds))

            progress_msg = f"Processed: {processed_count}, Skipped: {skipped_count}, Total: {total_processed}/{total_images} images. ETA: {eta_str}"
            print(progress_msg)  # Display on CMD
            yield progress_msg  # Update Gradio interface

            if total_processed >= total_images:
                break

    progress_bar.close()
    total_time = time.time() - start_time
    completion_msg = f"Face extraction complete! Processed: {processed_count}, Skipped: {skipped_count}, Total: {total_processed}/{total_images} images in {time.strftime('%H:%M:%S', time.gmtime(total_time))}"
    print(completion_msg)
    yield completion_msg

# --- Duplicate Detection ---

def calculate_color_hash(img, bins=32):
    """Calculates a color histogram hash for an image."""
    hist = cv2.calcHist([cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)], [0, 1, 2], None, [bins, bins, bins], [0, 256, 0, 256, 0, 256])
    hist = cv2.normalize(hist, hist).flatten()
    return hist

def compare_color_histograms(hist1, hist2, method=cv2.HISTCMP_CHISQR):
    """Compares two color histograms using the specified method."""
    return cv2.compareHist(hist1, hist2, method)

def extract_kaze_features(image):
    """Extracts KAZE features from an image."""
    kaze = cv2.KAZE_create()
    keypoints, descriptors = kaze.detectAndCompute(image, None)
    return keypoints, descriptors

def match_features(desc1, desc2, method='BF', ratio_test=0.75):
    """Matches features between two sets of descriptors."""
    if method == 'BF':
        matcher = cv2.BFMatcher()
        matches = matcher.knnMatch(desc1, desc2, k=2)
    elif method == 'FLANN':
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        matcher = cv2.FlannBasedMatcher(index_params, search_params)
        matches = matcher.knnMatch(desc1, desc2, k=2)
    else:
        raise ValueError("Invalid matching method. Choose 'BF' or 'FLANN'.")

    # Apply ratio test to filter good matches
    good_matches = []
    for m, n in matches:
        if m.distance < ratio_test * n.distance:
            good_matches.append(m)

    return good_matches

def process_image_duplicates(args):
    filename, folder_path, destination_folder, hash_algorithms, hash_size, cutoff, move_duplicates = args
    image_path = os.path.join(folder_path, filename)

    try:
        if isinstance(hash_algorithms, str):
            hash_algorithms = [hash_algorithms]

        with Image.open(image_path) as img:
            if img.mode != 'RGB':
                img = img.convert('RGB')

            image_gray = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY)
            image_hashes = {}
            duplicates_found = set()

            # Handle KAZE features
            if 'kaze' in hash_algorithms:
                try:
                    keypoints, descriptors = extract_kaze_features(image_gray)
                    serializable_keypoints = [(kp.pt, kp.size, kp.angle, kp.response, kp.octave, kp.class_id) for kp in keypoints]
                    image_hashes['kaze'] = (serializable_keypoints, descriptors)
                except Exception as e:
                    logging.error(f"Error processing {filename} with kaze: {str(e)}")

            # Calculate other hashes
            for alg in hash_algorithms:
                if alg == 'kaze':
                    continue
                try:
                    if alg == 'ahash':
                        hash_obj = imagehash.average_hash(img, hash_size=int(hash_size))
                    elif alg == 'dhash':
                        hash_obj = imagehash.dhash(img, hash_size=int(hash_size))
                    elif alg == 'phash':
                        hash_obj = imagehash.phash(img, hash_size=int(hash_size))
                    elif alg == 'whash':
                        hash_obj = imagehash.whash(img, hash_size=int(hash_size))
                    elif alg == 'colorhash':
                        hash_obj = calculate_color_hash(img)
                    image_hashes[alg] = hash_obj
                except Exception as e:
                    logging.error(f"Error processing {filename} with {alg}: {str(e)}")
                    continue

            if not image_hashes:
                return f"Failed to generate any hashes for {filename}", 0, 0, filename, []

            return "Processed: {}".format(filename), 0, 0, filename, image_hashes

    except Exception as e:
        logging.error(f"Error processing {filename}: {str(e)}")
        return f"Error processing {filename}: {str(e)}", 0, 0, filename, {}

def find_duplicates_multi_hash(folder_path, destination_folder, move_duplicates, hash_algorithms, hash_size, cutoff, num_threads):
    os.makedirs(destination_folder, exist_ok=True)

    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(img_formats)]
    total_files = len(image_files)
    processed_count = 0
    skipped_count = 0
    start_time = time.time()

    results_list = []

    with ProcessPoolExecutor(max_workers=num_threads) as executor:
        futures = [
            executor.submit(process_image_duplicates,
                          (f, folder_path, destination_folder, hash_algorithms,
                           hash_size, cutoff, move_duplicates))
            for f in image_files
        ]

        for future in as_completed(futures):
            try:
                result, processed, skipped, filename, image_hashes = future.result()
                results_list.append((filename, image_hashes))

                processed_count += 1

                elapsed_time = time.time() - start_time
                remaining_files = total_files - processed_count
                files_per_second = processed_count / elapsed_time if elapsed_time > 0 else 0
                eta_seconds = remaining_files / files_per_second if files_per_second > 0 else 0
                eta_str = time.strftime("%H:%M:%S", time.gmtime(eta_seconds))
                progress_msg = f"{result} (Processed: {processed_count}, Skipped: {skipped_count}, Total: {processed_count}/{total_files}, ETA: {eta_str})"
                print(progress_msg)
                yield progress_msg

            except Exception as exc:
                logging.error(f"Generated an exception: {exc}")

    # --- Merge and Move Duplicates (performed after collecting results) ---

    uf = UnionFind()
    all_relationships = []
    all_image_hashes = {}

    for filename, image_hashes in results_list:
        all_image_hashes[filename] = image_hashes
        for alg, current_hash in image_hashes.items():
            if alg not in ['kaze', 'colorhash']:
                current_hash = str(current_hash)

            for other_filename, other_hashes in all_image_hashes.items():
                if filename == other_filename:
                    continue

                if alg in other_hashes:
                    other_hash = other_hashes[alg]
                    if alg not in ['kaze', 'colorhash']:
                        other_hash = str(other_hash)

                    hash_cutoff = float(cutoff.get(alg, 30)) if alg == 'colorhash' else int(cutoff.get(alg, 2))

                    if alg == 'kaze':
                        try:
                            matches = match_features(image_hashes['kaze'][1], other_hashes['kaze'][1])
                            if len(matches) > int(cutoff.get(alg, 30)):
                                uf.union(filename, other_filename)
                                all_relationships.append((filename, other_filename, alg))
                        except Exception as e:
                            logging.error(f"Error comparing KAZE features for {filename} with {other_filename}: {str(e)}")
                    elif alg == 'colorhash':
                        try:
                            if compare_color_histograms(current_hash, other_hash) < hash_cutoff:
                                uf.union(filename, other_filename)
                                all_relationships.append((filename, other_filename, alg))
                        except Exception as e:
                            logging.error(f"Error comparing ColorHash for {filename} with {other_filename}: {str(e)}")
                    else:
                        if abs(int(current_hash, 16) - int(other_hash, 16)) < hash_cutoff:
                            uf.union(filename, other_filename)
                            all_relationships.append((filename, other_filename, alg))

    groups = {}
    for file in uf.parent:
        root = uf.find(file)
        if root not in groups:
            groups[root] = set()
        groups[root].add(file)

    duplicate_groups = [group for group in groups.values() if len(group) > 1]

    # Use a set to track moved files
    moved_files = set()

    if move_duplicates:
        moved_count = 0
        for group in duplicate_groups:
            # Convert set to list for consistent sorting
            group_list = list(group)

            # Find the "original" based on the file with the lexicographically smallest name
            original = min(group_list)

            # Move duplicates only if they haven't been moved already
            for duplicate_file in group_list:
                if duplicate_file != original and duplicate_file not in moved_files:
                    try:
                        source_path = os.path.join(folder_path, duplicate_file)
                        dest_path = os.path.join(destination_folder, duplicate_file)
                        if os.path.exists(source_path) and not os.path.exists(dest_path):
                            shutil.move(source_path, dest_path)
                            moved_files.add(duplicate_file)  # Add moved file to the set
                            moved_count += 1
                            print(f"Moved {duplicate_file} to destination folder (Original: {original})")
                    except Exception as e:
                        logging.error(f"Error moving duplicate {duplicate_file}: {str(e)}")

        print(f"Moved {moved_count} duplicate files to destination folder")
        yield f"Moved {moved_count} duplicate files to destination folder"

    # Save duplicates information (same as before)
    with open(os.path.join(destination_folder, 'duplicates.txt'), 'w') as f:
        f.write("Direct relationships found:\n")
        for rel in all_relationships:
            f.write(f"{rel[0]} is a duplicate of {rel[1]} (Method: {rel[2]})\n")

        f.write("\nFinal duplicate groups:\n")
        for group in duplicate_groups:
            sorted_group = sorted(list(group))
            original = sorted_group[0]
            duplicates = sorted_group[1:]
            f.write(f"\nGroup with original: {original}\n")
            for duplicate in duplicates:
                f.write(f"  - {duplicate} is a duplicate of {original}\n")

    total_duplicates = sum(len(group) - 1 for group in duplicate_groups)
    final_progress_msg = f"Found {total_duplicates} duplicates in {len(duplicate_groups)} groups. Results saved to duplicates.txt. Processed: {processed_count}, Skipped: {skipped_count}, Total: {processed_count}/{total_files}"
    print(final_progress_msg)
    yield final_progress_msg

# UnionFind class (for grouping duplicates)
class UnionFind:
    def __init__(self):
        self.parent = {}
        self.rank = {}

    def find(self, x):
        if x not in self.parent:
            self.parent[x] = x
            self.rank[x] = 0
            return x
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  # Path compression
        return self.parent[x]

    def union(self, x, y):
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            if self.rank[root_x] < self.rank[root_y]:
                root_x, root_y = root_y, root_x
            self.parent[root_y] = root_x
            if self.rank[root_x] == self.rank[root_y]:
                self.rank[root_x] += 1

def main():
    parser = argparse.ArgumentParser(description="Ultimate Image Processing Toolkit by SECourses")
    parser.add_argument("--share", action="store_true", help="Create a publicly shareable link")
    args = parser.parse_args()

    with gr.Blocks(title="Ultimate Image Processing Toolkit by SECourses") as demo:
        gr.Markdown("# Ultimate Image Processing Toolkit by SECourses V9 > https://www.patreon.com/posts/112126955")

        with gr.Tabs():
            with gr.Tab("Image Cropper"):
                gr.Markdown("Crop images using YOLO object detection or SAM2 segmentation")
                input_folder = gr.Textbox(label="Input Folder")
                output_folder = gr.Textbox(label="Output Folder")
                aspect_ratios = gr.Textbox(label="Aspect Ratios (comma-separated, e.g., 1024x1024,1280x720,1920x1080,...)", value="1024x1024")
                yolo_folder = gr.Textbox(label="YOLO/SAM2 Detected Images Folder (Optional)")
                save_yolo = gr.Checkbox(label="Save YOLO/SAM2 Detected Images")
                num_threads = gr.Slider(minimum=1, maximum=os.cpu_count(), step=1, label="Number of Threads", value=4)
                overwrite_crop = gr.Checkbox(label="Overwrite existing files", value=False)
                with gr.Row():
                    selected_class = gr.Dropdown(choices=YOLO_CLASSES, label="Select Detection Class (YOLO)", value="person")
                    sam2_prompt = gr.Textbox(label="SAM2 Prompt (Leave empty to use YOLO)", placeholder="Enter object to segment")
                save_as_png = gr.Checkbox(label="Save as PNG (default) or JPG", value=True)

                crop_button = gr.Button("Crop Images")
                crop_output = gr.Textbox(label="Progress", lines=5)

                crop_button.click(crop_images, inputs=[input_folder, output_folder, aspect_ratios, yolo_folder, save_yolo, num_threads, overwrite_crop, selected_class, save_as_png, sam2_prompt], outputs=crop_output)

            with gr.Tab("Image Resizer"):
                gr.Markdown("Resize images to various dimensions")
                resize_input_folder = gr.Textbox(label="Input Folder")
                resize_output_folder = gr.Textbox(label="Output Folder")
                resize_resolutions = gr.Textbox(label="Resolutions (comma-separated, e.g., 1024x1024,1280x720,1920x1080,...)", value="1024x1024")
                save_png_resize = gr.Checkbox(label="Save as PNG (default) or JPG", value=True)
                num_threads_resize = gr.Slider(minimum=1, maximum=os.cpu_count(), step=1, label="Number of Threads", value=4)
                overwrite_resize = gr.Checkbox(label="Overwrite existing files", value=False)
                no_crop_resize = gr.Checkbox(label="Resize without cropping (center image and add white background)", value=False)
                resize_button = gr.Button("Resize Images")
                resize_output = gr.Textbox(label="Progress", lines=5)

                resize_button.click(resize_images, inputs=[resize_input_folder, resize_output_folder, resize_resolutions, save_png_resize, num_threads_resize, overwrite_resize, no_crop_resize], outputs=resize_output)

            with gr.Tab("Move Low-Res Files"):
                gr.Markdown("Move low-resolution files between folders")
                folder_a = gr.Textbox(label="Folder A (Cropped images)")
                folder_b = gr.Textbox(label="Folder B (Raw images)")
                folder_c = gr.Textbox(label="Folder C (Raw low-res images)")
                folder_d = gr.Textbox(label="Folder D (Cropped low-res images)")
                min_width = gr.Number(label="Minimum Width", value=1536)
                min_height = gr.Number(label="Minimum Height", value=1536)
                move_button = gr.Button("Move Low-Res Files")
                move_output = gr.Textbox(label="Output")

                move_button.click(move_low_res_files, inputs=[folder_a, folder_b, folder_c, folder_d, min_width, min_height], outputs=move_output)

            with gr.Tab("Rename Files"):
                gr.Markdown("Rename files in a folder")
                rename_folder = gr.Textbox(label="Folder Path")
                start_number = gr.Number(label="Starting Number", value=10001)
                rename_button = gr.Button("Rename Files")
                rename_output = gr.Textbox(label="Output")

                rename_button.click(rename_files, inputs=[rename_folder, start_number], outputs=rename_output)

            with gr.Tab("Extract Faces"):
                gr.Markdown("Extract faces from images")
                extract_input_folder = gr.Textbox(label="Input Folder")
                extract_output_folder = gr.Textbox(label="Output Folder")
                padding = gr.Number(label="Padding (pixels)", value=0,
                                    info="Expand or decrease face area. Positive values expand, negative values decrease.")
                save_png_extract = gr.Checkbox(label="Save as PNG (default) or JPG", value=True)
                num_threads_extract = gr.Slider(minimum=1, maximum=os.cpu_count(), step=1, label="Number of Threads", value=4)
                extract_button = gr.Button("Extract Faces")
                extract_output = gr.Textbox(label="Progress", lines=5)

                extract_button.click(extract_faces,
                                     inputs=[extract_input_folder, extract_output_folder, padding, save_png_extract, num_threads_extract],
                                     outputs=extract_output)

            with gr.Tab("Find Duplicates"):
                gr.Markdown("Find duplicate images using perceptual hashing and other methods")

                with gr.Row():
                    with gr.Column():
                        dup_source_folder = gr.Textbox(label="Source Folder")
                        dup_dest_folder = gr.Textbox(label="Destination Folder")

                        with gr.Row():
                            move_duplicates = gr.Checkbox(label="Move duplicates to destination folder", value=True)
                            hash_algorithm = gr.Dropdown(
                                choices=['ahash', 'dhash', 'phash', 'whash', 'colorhash', 'kaze', 'multi-hash'],
                                label="Hashing Algorithm",
                                value='phash'  # Default to multi-hash
                            )

                        default_hash_size = gr.Slider(minimum=4, maximum=32, step=1, label="Hash Size", value=16) # Increased default hash size

                        with gr.Accordion("Algorithm Details and Parameters", open=False):
                            gr.Markdown("""
                                **Algorithm Details:**

                                - **aHash (Average Hash):**
                                    - **How it works:** Calculates the average color of the image, then creates a hash based on whether each pixel is brighter or darker than the average.
                                    - **Pros:** Fast and simple. Good for detecting exact duplicates or very similar images.
                                    - **Cons:** Sensitive to small changes in color or brightness.
                                    - **Default Cutoff:** 2 (Lower = stricter, detects more duplicates)
                                    - **Example:** Imagine two identical photos, but one is slightly brighter. aHash might miss this as a duplicate if the brightness difference is significant enough.

                                - **dHash (Difference Hash):**
                                    - **How it works:** Similar to aHash but compares the difference in brightness between adjacent pixels.
                                    - **Pros:** Relatively fast, better than aHash at detecting small color changes.
                                    - **Cons:** Still can be sensitive to minor image alterations.
                                    - **Default Cutoff:** 5 (Lower = stricter)
                                    - **Example:** Two photos with slightly different color grading might be detected as duplicates by dHash.

                                - **pHash (Perceptual Hash):**
                                    - **How it works:** Uses the Discrete Cosine Transform (DCT) to focus on low-frequency components, making it robust to minor changes.
                                    - **Pros:** Very good at detecting duplicates even with compression, resizing, or minor edits.
                                    - **Cons:** Slower than aHash or dHash.
                                    - **Default Cutoff:** 8 (Lower = stricter)
                                    - **Example:** A compressed version of an original image is very likely to be identified as a duplicate by pHash.

                                - **wHash (Wavelet Hash):**
                                    - **How it works:** Uses Wavelet transform to capture image details at different scales.
                                    - **Pros:** Robust to resizing and minor color adjustments.
                                    - **Cons:** Can be slower than other hash methods.
                                    - **Default Cutoff:** 3 (Lower = stricter)
                                    - **Example:** wHash is good at finding duplicates where one image has been slightly cropped or resized.

                                - **ColorHash:**
                                    - **How it works:** Creates a hash based on the color histogram of the image.
                                    - **Pros:** Detects images with similar color distributions.
                                    - **Cons:** Sensitive to changes in image content that maintain the same color palette.
                                    - **Default Cutoff:** 50 (Higher = stricter, detects fewer duplicates)
                                    - **Example:** Two different images with a similar color scheme (e.g., both predominantly blue) might be flagged as duplicates by ColorHash.

                                - **KAZE:**
                                    - **How it works:** Detects and describes local features in images using non-linear scale spaces, making it robust to scale, rotation, and some level of viewpoint change.
                                    - **Pros:** Excellent at finding near-duplicates with significant changes in perspective, scale, or rotation. Generally more robust and accurate than SIFT/SURF.
                                    - **Cons:** Computationally more expensive and slower than hashing methods, but typically faster or comparable to SIFT/SURF. Requires the 'opencv-contrib-python' package.
                                    - **Default Matches:** 30 (Higher = stricter, detects fewer duplicates)
                                    - **Example:** KAZE is well-suited for detecting duplicates in images that have undergone various transformations like resizing, rotation, or moderate viewpoint changes.

                                - **Multi-Hash:**
                                    - **How it works:** Combines multiple hashing algorithms (aHash, dHash, pHash, wHash, ColorHash, KAZE) to improve overall accuracy.
                                    - **Pros:** Leverages the strengths of different methods, leading to better duplicate detection.
                                    - **Cons:** Can be slower as it performs multiple calculations.
                                    - **Example:** By combining the results of different algorithms, Multi-Hash can catch a wider range of duplicates that might be missed by individual methods.
                                """)

                        with gr.Row():
                            default_cutoff_ahash = gr.Slider(minimum=0, maximum=10, step=1, label="aHash Cutoff", value=5)
                            default_cutoff_dhash = gr.Slider(minimum=0, maximum=10, step=1, label="dHash Cutoff", value=5)
                            default_cutoff_phash = gr.Slider(minimum=0, maximum=10, step=1, label="pHash Cutoff", value=8)

                        with gr.Row():
                            default_cutoff_whash = gr.Slider(minimum=0, maximum=10, step=1, label="wHash Cutoff", value=3)
                            default_cutoff_colorhash = gr.Slider(minimum=0, maximum=100, step=5, label="ColorHash Cutoff", value=50)
                            default_cutoff_kaze = gr.Slider(minimum=0, maximum=100, step=5, label="KAZE Matches", value=40)

                        num_threads_duplicates = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Threads", value=4)

                        find_dup_button = gr.Button("Find Duplicates")
                        find_dup_output = gr.Textbox(label="Output", lines=5)

                def find_duplicates_wrapper(*args):
                    try:
                        dup_source_folder, dup_dest_folder, move_duplicates, hash_algorithm, default_hash_size, \
                        default_cutoff_ahash, default_cutoff_dhash, default_cutoff_phash, default_cutoff_whash, \
                        default_cutoff_colorhash, default_cutoff_kaze, num_threads_duplicates = args

                        # Validate input folders
                        if not os.path.exists(dup_source_folder):
                            return f"Source folder does not exist: {dup_source_folder}"

                        # Create destination folder if it doesn't exist
                        os.makedirs(dup_dest_folder, exist_ok=True)

                        hash_size = int(default_hash_size)

                        # Always convert to list of algorithms
                        if hash_algorithm == 'multi-hash':
                            selected_algorithms = ['ahash', 'dhash', 'phash', 'whash', 'colorhash', 'kaze']
                        else:
                            selected_algorithms = [hash_algorithm]

                        # Initialize cutoff values for selected algorithms
                        cutoff = {
                            'ahash': int(default_cutoff_ahash),
                            'dhash': int(default_cutoff_dhash),
                            'phash': int(default_cutoff_phash),
                            'whash': int(default_cutoff_whash),
                            'colorhash': float(default_cutoff_colorhash),
                            'kaze': int(default_cutoff_kaze)
                        }

                        # Initialize global dictionaries for selected algorithms
                        global global_hash_dicts
                        global_hash_dicts = {alg: {} for alg in selected_algorithms}

                        # Reset duplicate images list
                        global global_duplicate_images
                        global_duplicate_images = []

                        output_text = ""
                        for message in find_duplicates_multi_hash(
                            dup_source_folder,
                            dup_dest_folder,
                            move_duplicates,
                            selected_algorithms,
                            hash_size,
                            cutoff,
                            num_threads_duplicates
                        ):
                            output_text += str(message) + "\n"

                        return output_text
                    except Exception as e:
                        logging.error(f"Error in duplicate detection: {str(e)}")
                        return f"Error in duplicate detection: {str(e)}"

                find_dup_button.click(
                    fn=find_duplicates_wrapper,
                    inputs=[
                        dup_source_folder,
                        dup_dest_folder,
                        move_duplicates,
                        hash_algorithm,
                        default_hash_size,
                        default_cutoff_ahash,
                        default_cutoff_dhash,
                        default_cutoff_phash,
                        default_cutoff_whash,
                        default_cutoff_colorhash,
                        default_cutoff_kaze,
                        num_threads_duplicates
                    ],
                    outputs=find_dup_output
                )
    demo.launch(share=args.share, inbrowser=True)

if __name__ == "__main__":
    main()