import streamlit as st
import face_recognition
import cv2
import numpy as np
import os
from PIL import Image
import dlib

# Check if CUDA is available
if dlib.DLIB_USE_CUDA:
    st.write("CUDA is available. Using GPU for face recognition.")
else:
    st.write("CUDA is not available. Using CPU for face recognition.")

# 📂 Ruta donde están las imágenes
IMAGE_DIRECTORY = "dataset_faces/"

# 📌 Función para cargar imágenes y extraer embeddings
def load_images_and_encodings(directory):
    known_encodings = []
    known_images = []
    known_names = []

    for filename in os.listdir(directory):
        if filename.endswith((".jpg", ".png", ".jpeg")):
            path = os.path.join(directory, filename)
            image = face_recognition.load_image_file(path)
            encodings = face_recognition.face_encodings(image)
            
            if encodings:  # Si encontró una cara
                known_encodings.append(encodings[0])
                known_images.append(path)
                known_names.append(filename)

    return known_encodings, known_images, known_names

# 📌 Cargar imágenes y sus embeddings al inicio
known_encodings, known_images, known_names = load_images_and_encodings(IMAGE_DIRECTORY)

# 📌 Función para encontrar imágenes similares
def find_similar_faces(uploaded_image):
    # Convertir a array de NumPy
    image_np = np.array(uploaded_image)
    face_encodings = face_recognition.face_encodings(image_np)

    if not face_encodings:
        return [], []  # Si no encuentra caras en la imagen subida

    query_encoding = face_encodings[0]
    distances = face_recognition.face_distance(known_encodings, query_encoding)
    
    # Ordenar las imágenes por similitud
    sorted_indices = np.argsort(distances)
    
    return [known_images[i] for i in sorted_indices], [distances[i] for i in sorted_indices]

# 📌 Interfaz en Streamlit
st.title("🔍 Buscador de Rostros en un Directorio")
st.write("Sube una imagen y te mostraremos las fotos más similares en el directorio.")

uploaded_file = st.file_uploader("📤 Sube una imagen", type=["jpg", "png", "jpeg"])

if uploaded_file:
    # Mostrar imagen subida
    uploaded_image = Image.open(uploaded_file)
    st.image(uploaded_image, caption="Imagen subida", use_container_width=True)
    
    # Buscar rostros similares
    similar_images, distances = find_similar_faces(uploaded_image)

    # Mostrar resultados
    if similar_images:
        st.subheader("📸 Imágenes más similares:")
        cols = st.columns(len(similar_images[:]))  # Crear columnas para las imágenes
        for idx, (col, image_path) in enumerate(zip(cols, similar_images[:])):  # Mostrar las 5 más similares
            col.image(image_path, caption=f"Similitud: {1 - distances[idx]:.2f}", use_container_width=True)
    else:
        st.warning("⚠ No se detectó ningún rostro en la imagen subida.")