import torch
import os
from fastapi import FastAPI, File, UploadFile, HTTPException
import cv2
from PIL import Image
import numpy as np
from io import BytesIO

# Configurar la ubicación del caché de PyTorch
os.environ["TORCH_HOME"] = "/tmp/torch_cache"  # Directorio accesible para almacenar el caché

# Iniciar la aplicación FastAPI
app = FastAPI()

# Cargar el modelo preentrenado YOLOv5 desde el repositorio oficial
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # 'yolov5s' es una versión ligera

# Función para buscar diferentes tipos de objetos según los modelos Haar
def buscar_existe(image):
    existe = "No detectado"
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Lista de modelos Haar para ojos, ojos con gafas y sonrisas
    models = [
        ('haarcascade_eye.xml', 'ojos'),
        ('haarcascade_eye_tree_eyeglasses.xml', 'ojos detrás de gafas'),
        ('haarcascade_smile.xml', 'sonrisa')
    ]
    
    # Probar los modelos Haar
    for model_file, descripcion in models:
        print(f"Probando modelo Haar: {model_file}")
        cascade = cv2.CascadeClassifier(model_file)
        if cascade.empty():
            print(f"El modelo {model_file} no se pudo cargar.")
            continue
        detecciones = cascade.detectMultiScale(gray, 1.3, 5, minSize=(10, 10))
        if len(detecciones) > 0:
            # Filtro adicional para "sonrisa"
            if descripcion != "sonrisa":
                existe = f"Detectado: {descripcion}"
                break
            else:
                existe = "No se detectó sonrisa, pero otras características podrían haber activado el modelo."
                continue

    # Usar el modelo YOLOv5 localmente para detectar objetos como sillas
    if existe == "No detectado":
        # Convertir la imagen a formato que YOLOv5 pueda recibir
        pil_image = Image.fromarray(image)
        pil_image = pil_image.convert("RGB")  # Convertir a RGB si es necesario
        
        # Realizar la predicción con YOLOv5
        results = model(pil_image)  # Ejecutar la predicción
        
        # Convertir los resultados a pandas DataFrame
        detected_objects = results.pandas().xywh[0]  # Devolver las predicciones en un DataFrame
        
        print(detected_objects)  # Para ver qué objetos están siendo detectados
        
        # Verificar si se detectaron sillas en los resultados
        if 'chair' in detected_objects['name'].values:
            existe = "Detectado: silla"
        else:
            existe = "No se detectó silla."
        
    return existe

# Ruta de predicción
@app.post('/predict/')
async def predict(file: UploadFile = File(...)):
    try:
        # Leer la imagen y convertirla a un array de numpy
        image = Image.open(BytesIO(await file.read()))
        image = np.asarray(image)

        # Buscar cualquier detección de silla
        prediction = buscar_existe(image)
        return {"prediction": prediction}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))