import cv2
import numpy as np
import torch
from PIL import Image
from facenet_pytorch import MTCNN, InceptionResnetV1
import gradio as gr
import zipfile
import os
import shutil
import traceback

# Load pre-trained model for face detection and face recognition
try:
    mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20)
    resnet = InceptionResnetV1(pretrained='vggface2').eval()
except Exception as e:
    print(f"Error loading models: {e}")

# Function to preprocess image and extract face embeddings
def preprocess_image(image):
    try:
        image = Image.fromarray(image)
        image_cropped = mtcnn(image)
        if image_cropped is not None:
            image_embedding = resnet(image_cropped.unsqueeze(0))
            return image_embedding
        else:
            print("No face detected in image during preprocessing.")
    except Exception as e:
        print(f"Error preprocessing image: {e}")
    return None

# Load embeddings from a temporary uploaded folder
def load_embeddings_from_folder(folder_path):
    embeddings = []
    try:
        for filename in os.listdir(folder_path):
            img_path = os.path.join(folder_path, filename)
            image = cv2.imread(img_path)
            if image is None:
                print(f"Image {filename} could not be read or is not in a supported format.")
                continue
            embedding = preprocess_image(image)
            if embedding is not None:
                embeddings.append((filename, embedding))
            else:
                print(f"Embedding for {filename} could not be generated.")
    except Exception as e:
        print(f"Error loading embeddings from folder {folder_path}: {e}")
    return embeddings

# Compare two embeddings and return True if they match
def compare_embeddings(embedding1, embedding2, threshold=0.3):
    try:
        similarity = torch.nn.functional.cosine_similarity(embedding1, embedding2)
        return similarity.item() > threshold
    except Exception as e:
        print(f"Error comparing embeddings: {e}")
    return False

# Identify images in Folder A that do not match any images in Folder B
def identify_non_matching_faces(folder_a_zip, folder_b_zip):
    try:
        # Create temporary directories
        folder_a_path = "temp_folder_a"
        folder_b_path = "temp_folder_b"
        non_matching_folder = "non_matching_images"
        os.makedirs(folder_a_path, exist_ok=True)
        os.makedirs(folder_b_path, exist_ok=True)
        os.makedirs(non_matching_folder, exist_ok=True)

        # Unzip the folders
        with zipfile.ZipFile(folder_a_zip, 'r') as zip_ref:
            zip_ref.extractall(folder_a_path)
        with zipfile.ZipFile(folder_b_zip, 'r') as zip_ref:
            zip_ref.extractall(folder_b_path)

        # Process embeddings
        embeddings_a = load_embeddings_from_folder(folder_a_path)
        embeddings_b = load_embeddings_from_folder(folder_b_path)

        if not embeddings_a:
            print("No valid embeddings found in Folder A.")
        if not embeddings_b:
            print("No valid embeddings found in Folder B.")

        # Track non-matching images
        non_matching_images = []

        for name_a, embedding_a in embeddings_a:
            match_found = False
            for name_b, embedding_b in embeddings_b:
                similarity = torch.nn.functional.cosine_similarity(embedding_a, embedding_b).item()
                print(f"Comparing {name_a} with {name_b} | Similarity: {similarity}")  # Debug: Log similarity score

                if similarity > 0.3:  # Adjust threshold if needed
                    match_found = True
                    break
            if not match_found:
                print(f"Non-matching image found: {name_a}")
                non_matching_images.append(name_a)
                shutil.copy(os.path.join(folder_a_path, name_a), os.path.join(non_matching_folder, name_a))

        if non_matching_images:
            # Create a zip file of all non-matching images
            zip_filename = "non_matching_images.zip"
            with zipfile.ZipFile(zip_filename, 'w') as zipf:
                for root, _, files in os.walk(non_matching_folder):
                    for file in files:
                        zipf.write(os.path.join(root, file), arcname=file)
            print("Non-matching images found and zipped.")
            return zip_filename  # Return the zip file path for download
        else:
            print("No non-matching images found.")
            shutil.rmtree(non_matching_folder)  # Clean up empty folder if no images
            return None

    except Exception as e:
        # Catch all other exceptions and return the error as a string
        error_message = f"An error occurred: {str(e)}\n{traceback.format_exc()}"
        print(error_message)
        return error_message

# Gradio interface
def gradio_interface(folder_a_zip, folder_b_zip):
    result = identify_non_matching_faces(folder_a_zip.name, folder_b_zip.name)
    if isinstance(result, str) and result.endswith('.zip'):
        return result  # Return zip file path for download
    elif result is None:
        return "No non-matching images found."  # Return a message if no non-matching images
    else:
        return result  # Return the error message

with gr.Blocks() as demo:
    gr.Markdown("# Face Comparison between Two Zipped Image Folders")

    with gr.Row():
        folder_a_input = gr.File(label="Upload Zipped Folder A", type="filepath")
        folder_b_input = gr.File(label="Upload Zipped Folder B", type="filepath")
        compare_faces_button = gr.Button("Find Non-Matching Faces")
        download_link = gr.File(label="Download Non-Matching Images Zip")

    compare_faces_button.click(gradio_interface, inputs=[folder_a_input, folder_b_input], outputs=[download_link])

# Launch the Gradio interface
demo.launch()