from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from keras.models import load_model
from tensorflow.keras.applications.inception_v3 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
import numpy as np
from PIL import Image
import io
import base64

app = FastAPI()

# Lista de clases
class_names = ['acanthoica', 'akashiwo', 'alexandrium', 'amoeba', 'amphidinium', 'amylax', 'apedinella', 
               'asterionellopsis', 'bacillaria', 'bacteriastrum', 'biddulphia', 'calciopappus', 'cerataulina', 
               'ceratium', 'chaetoceros', 'chrysochromulina', 'cochlodinium', 'corethron', 'corymbellus', 
               'coscinodiscus', 'cryptophyta', 'cylindrotheca', 'dactyliosolen', 'delphineis', 'dictyocha', 
               'dinobryon', 'dinophysis', 'ditylum', 'emiliania', 'ephemera', 'eucampia', 'euglena', 
               'gonyaulax', 'guinardia', 'gyrodinium', 'hemiaulus', 'heterocapsa', 'karenia', 'katodinium', 
               'kryptoperidinium', 'laboea', 'lauderia', 'leptocylindrus', 'licmophora', 'nanoneis', 
               'odontella', 'ophiaster', 'ostreopsis', 'oxytoxum', 'paralia', 'parvicorbicula', 'phaeocystis', 
               'pleuronema', 'pleurosigma', 'polykrikos', 'prorocentrum', 'proterythropsis', 'protoperidinium', 
               'pseudo-nitzschia', 'pseudochattonella', 'pyramimonas', 'rhabdolithes', 'rhizosolenia', 
               'scrippsiella', 'skeletonema', 'stephanopyxis', 'syracosphaera', 'thalassionema', 'thalassiosira', 
               'trichodesmium', 'vicicitus', 'warnowia']

# Descargar y cargar el modelo desde Hugging Face Hub
model_path = hf_hub_download(repo_id="Daniel00611/InceptionV3_72", filename="InceptionV3_72.keras")
model = load_model(model_path)

def preprocess_image(img, target_size=(299, 299)):
    # Convertir a RGB si la imagen no está en ese formato
    if img.mode != "RGB":
        img = img.convert("RGB")
    img = img.resize(target_size)
    img_array = img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    return img_array

# Modelo para recibir múltiples imágenes en Base64
class ImagesBase64(BaseModel):
    images_base64: list[str]  # Lista de imágenes en formato Base64

# Ruta para imágenes subidas como archivo
@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    try:
        # Leer la imagen subida
        img = Image.open(io.BytesIO(await file.read()))
        img_array = preprocess_image(img)
        
        # Realizar predicción
        predictions = model.predict(img_array)[0]
        
        # Obtener el top 10 de predicciones
        top_10_indices = predictions.argsort()[-10:][::-1]
        top_10_classes = [class_names[i] for i in top_10_indices]
        top_10_probabilities = predictions[top_10_indices]
        
        # Formar respuesta en formato JSON
        result = [{"class": top_10_classes[i], "probability": float(top_10_probabilities[i])} for i in range(10)]
        return {"predictions": result}
    
    except Exception as e:
        return {"error": str(e)}

# Ruta para imágenes en formato Base64
@app.post("/predict_base64/")
async def predict_base64(image_data: ImagesBase64):
    results = {}
    try:
        for index, image_base64 in enumerate(image_data.images_base64):
            # Decodificar cada imagen Base64
            image_bytes = base64.b64decode(image_base64)
            img = Image.open(io.BytesIO(image_bytes))
            img_array = preprocess_image(img)
            
            # Realizar predicción
            predictions = model.predict(img_array)[0]
            
            # Obtener el top 10 de predicciones
            top_10_indices = predictions.argsort()[-10:][::-1]
            top_10_classes = [class_names[i] for i in top_10_indices]
            top_10_probabilities = predictions[top_10_indices]
            
            # Formar respuesta para la imagen actual
            image_result = [{"class": top_10_classes[i], "probability": float(top_10_probabilities[i])} for i in range(10)]
            results[f"imagen{index + 1}"] = image_result
        
        return results
    
    except Exception as e:
        return {"error": str(e)}
@app.get("/")
def greet_json():
    return {"Hello": "World!"}