from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import tensorflow as tf
import numpy as np

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

# Cargar el modelo
model = tf.keras.models.load_model("model.h5")

# Clase para los datos de entrada
class PredictionInput(BaseModel):
    pclass: int
    gender: str  # "male" o "female"
    age: float
    sibsp: int
    parch: int
    fare: float

# Función de preprocesamiento
def preprocess_input(data: PredictionInput):
    # Convertir género a un valor numérico
    gender = 1 if data.gender.lower() == "female" else 0
    
    # Crear array con los datos procesados
    input_data = np.array([[data.pclass, gender, data.age, data.sibsp, data.parch, data.fare]])
    
    return input_data

# Endpoint de predicción
@app.post("/predict")
async def predict(data: PredictionInput):
    try:
        # Preprocesar los datos de entrada
        input_data = preprocess_input(data)
        
        # Realizar la predicción
        prediction = model.predict(input_data)
        
        # Devolver el resultado en un diccionario
        return {"survived_probability": float(prediction[0][0])}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))