import os
import numpy as np
from PIL import Image
import gradio as gr
from deepface import DeepFace
from datasets import load_dataset

# Cargamos el dataset sin aplicar formato especial
dataset = load_dataset("Segizu/dataset_faces")

# Cargar embeddings de todas las imágenes del dataset
def build_database():
    database = []
    for filename in os.listdir(IMAGE_DIRECTORY):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            path = os.path.join(IMAGE_DIRECTORY, filename)
            try:
                representation = DeepFace.represent(img_path=path, model_name="Facenet")[0]["embedding"]
                database.append((filename, path, representation))
            except:
                print(f"❌ No se pudo procesar: {filename}")
    return database

# Inicializamos base de datos
database = build_database()

# Comparar imagen cargada con las del dataset
def find_similar_faces(uploaded_image):
    try:
        uploaded_image = np.array(uploaded_image)
        query_representation = DeepFace.represent(img_path=uploaded_image, model_name="Facenet")[0]["embedding"]
    except:
        return [], "⚠ No se detectó un rostro válido en la imagen."

    similarities = []
    for name, path, rep in database:
        distance = np.linalg.norm(np.array(query_representation) - np.array(rep))
        similarity = 1 / (1 + distance)  # Normalizamos para que 1 = muy similar
        similarities.append((similarity, name, path))

    # Ordenar por similitud
    similarities.sort(reverse=True)
    top_matches = similarities[:5]

    # Formatear salida para gradio
    gallery_items = []
    text_summary = ""
    for sim, name, path in top_matches:
        img = Image.open(path)
        caption = f"{name} - Similitud: {sim:.2f}"
        gallery_items.append({"image": img, "caption": caption})
        text_summary += caption + "\n"

    return gallery_items, text_summary

# Interfaz Gradio
demo = gr.Interface(
    fn=find_similar_faces,
    inputs=gr.Image(label="📤 Sube una imagen", type="pil"),
    outputs=[
        gr.Gallery(label="📸 Rostros más similares"),
        gr.Textbox(label="🧠 Similitud", lines=6)
    ],
    title="🔍 Buscador de Rostros con DeepFace",
    description="Sube una imagen y te mostrará los rostros más similares desde el directorio `dataset_faces/`."
)

demo.launch()