import numpy as np
from PIL import Image
import gradio as gr
from deepface import DeepFace
from datasets import load_dataset, DownloadConfig
import os
os.system("rm -rf ~/.cache/huggingface/hub/datasets--Segizu--dataset_faces")

# ✅ Cargar el dataset de Hugging Face forzando la descarga limpia
download_config = DownloadConfig(force_download=True)
dataset = load_dataset("Segizu/dataset_faces", download_config=download_config)
if "train" in dataset:
    dataset = dataset["train"]

# 🔄 Preprocesar imagen para Facenet
def preprocess_image(img):
    img_rgb = img.convert("RGB")
    img_resized = img_rgb.resize((160, 160), Image.Resampling.LANCZOS)
    return np.array(img_resized)

# 📦 Construir base de datos de embeddings
def build_database():
    database = []
    for i, item in enumerate(dataset):
        try:
            img = item["image"]
            img_processed = preprocess_image(img)
            embedding = DeepFace.represent(
                img_path=img_processed,
                model_name="Facenet",
                enforce_detection=False
            )[0]["embedding"]
            database.append((f"image_{i}", img, embedding))
        except Exception as e:
            print(f"❌ No se pudo procesar imagen {i}: {e}")
    return database

# 🔍 Buscar rostros similares
def find_similar_faces(uploaded_image):
    try:
        img_processed = preprocess_image(uploaded_image)
        query_embedding = DeepFace.represent(
            img_path=img_processed,
            model_name="Facenet",
            enforce_detection=False
        )[0]["embedding"]
    except:
        return [], "⚠ No se detectó un rostro válido en la imagen."

    similarities = []
    for name, db_img, embedding in database:
        dist = np.linalg.norm(np.array(query_embedding) - np.array(embedding))
        sim_score = 1 / (1 + dist)
        similarities.append((sim_score, name, db_img))

    similarities.sort(reverse=True)
    top_matches = similarities[:]

    gallery_items = []
    text_summary = ""
    for sim, name, img in top_matches:
        caption = f"{name} - Similitud: {sim:.2f}"
        gallery_items.append((img, caption))
        text_summary += caption + "\n"

    return gallery_items, text_summary

# ⚙️ Inicializar base
database = build_database()

# 🎛️ 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 se comparará contra los rostros del dataset alojado en Hugging Face (`Segizu/dataset_faces`)."
)

demo.launch()