import gradio as gr
import face_recognition
import numpy as np
import os
from PIL import Image
import dlib

# Verificar si CUDA está disponible y seleccionar el modelo adecuado
if dlib.DLIB_USE_CUDA:
    print("✅ CUDA está disponible. Se usará GPU para reconocimiento facial.")
    model_used = "cnn"  # Modelo optimizado para GPU
else:
    print("⚠ CUDA no está disponible. Se usará CPU para reconocimiento facial.")
    model_used = "hog"  # Modelo más adecuado para CPU

# 📂 Directorio donde se encuentran las imágenes
IMAGE_DIRECTORY = "dataset_faces/"

def load_images_and_encodings(directory):
    """
    Carga las imágenes y extrae sus embeddings.
    """
    known_encodings = []
    known_images = []
    known_names = []

    for filename in os.listdir(directory):
        if filename.lower().endswith((".jpg", ".png", ".jpeg")):
            path = os.path.join(directory, filename)
            image = face_recognition.load_image_file(path)
            encodings = face_recognition.face_encodings(image, model=model_used)
            if encodings:  # Si se detecta al menos una cara
                known_encodings.append(encodings[0])
                known_images.append(path)
                known_names.append(filename)
    return known_encodings, known_images, known_names

# Cargar los datos de la carpeta de imágenes
known_encodings, known_images, known_names = load_images_and_encodings(IMAGE_DIRECTORY)

def find_similar_faces_gradio(uploaded_image):
    """
    Dada una imagen subida, busca las imágenes similares del dataset.
    Devuelve una lista de diccionarios para la galería y un texto con detalles.
    """
    if uploaded_image is None:
        return [], "No se subió ninguna imagen."
    
    # Convertir la imagen subida a array de NumPy
    image_np = np.array(uploaded_image)
    face_encodings = face_recognition.face_encodings(image_np, model=model_used)
    if not face_encodings:
        return [], "⚠ No se detectó ningún rostro en la imagen subida."
    
    query_encoding = face_encodings[0]
    distances = face_recognition.face_distance(known_encodings, query_encoding)
    sorted_indices = np.argsort(distances)  # Ordenar por similitud (menor distancia = mayor similitud)
    
    # Mostrar las 5 imágenes más similares
    top_n = 5
    gallery_items = []
    details = ""
    for idx in sorted_indices[:top_n]:
        # Abrir la imagen del dataset
        img = Image.open(known_images[idx])
        similarity = 1 - distances[idx]  # Definir similitud (valor entre 0 y 1)
        caption = f"{os.path.basename(known_images[idx])}: Similitud: {similarity:.2f}"
        gallery_items.append({"image": img, "caption": caption})
        details += caption + "\n"
    
    return gallery_items, details

# Definir la interfaz con Gradio
demo = gr.Interface(
    fn=find_similar_faces_gradio,
    inputs=gr.Image(label="Sube una imagen", type="pil"),
    outputs=[
        gr.Gallery(label="Imágenes similares").style(grid=[2], height="auto"),
        gr.Textbox(label="Detalles de similitud", lines=5)
    ],
    title="🔍 Buscador de Rostros en un Directorio",
    description="Sube una imagen y se mostrarán las fotos más similares del directorio."
)

demo.launch()