from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import numpy as np
import joblib

# Definir la clase InputData usando Pydantic
class InputData(BaseModel):
    age: float
    sex: float
    cp: float
    trestbps: float
    chol: float
    fbs: float
    restecg: float
    thalach: float
    exang: float
    oldpeak: float
    slope: float
    ca: float
    thal: float

app = FastAPI()

# Función para cargar el modelo
def load_model():
    model = joblib.load("modelo_heartdisease.pkl")
    return model

model = load_model()

@app.get("/")
async def root():
    return {"message": "API funcionando. Usa /predict/ para hacer predicciones."}

@app.post("/predict/")
async def predict(data: InputData):
    try:
        input_data = np.array([
            data.age,
            data.sex,
            data.cp,
            data.trestbps,
            data.chol,
            data.fbs,
            data.restecg,
            data.thalach,
            data.exang,
            data.oldpeak,
            data.slope,
            data.ca,
            data.thal
        ]).reshape(1, -1)

        prediction = model.predict(input_data)
        prediction_label = "Presencia de enfermedad cardíaca" if prediction[0] == 1 else "Ausencia de enfermedad cardíaca"
        return {"prediction": prediction_label}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))