import os
from io import BytesIO
from typing import Dict

import uvicorn
from fastapi import FastAPI, File, UploadFile
from PIL import Image
from pydantic import BaseModel
import tensorflow as tf

# Carga el modelo SavedModel (ajusta la ruta si es necesario)
model = tf.saved_model.load("./efficientnet_alzheimer")

app = FastAPI(title="API de clasificación de Alzheimer")

# Define la función de preprocesamiento de imágenes (la misma que en inference.py)
def preprocess_image(image):
    image = image.resize((200, 200))
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    return image

# Define el formato de la respuesta de la API
class Prediction(BaseModel):
    prediction: str
    confidence: float

# Endpoint para la predicción
@app.post("/predict", response_model=Prediction)
async def predict_image(image: UploadFile = File(...)):
    # Lee la imagen
    contents = await image.read()
    image = Image.open(BytesIO(contents))

    # Preprocesa la imagen
    processed_image = preprocess_image(image)

    # Realiza la predicción
    predictions = model(processed_image)

    # Obtén la clase predicha y la probabilidad
    predicted_class = np.argmax(predictions)
    confidence = np.max(predictions)

    # Mapea la clase numérica al nombre de la clase (igual que en inference.py)
    class_names = {
        0: "VeryMildDemented",
        1: "NonDemented",
        2: "ModerateDemented",
        3: "MildDemented"
    }
    predicted_class_name = class_names.get(predicted_class, "Clase desconocida")

    # Devuelve la predicción y la confianza
    return {"prediction": predicted_class_name, "confidence": confidence.item()}

# Punto de entrada para ejecutar la aplicación con Uvicorn
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8000)))