from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer

app = FastAPI()

# Cargar el modelo preentrenado desde el archivo .h5
model = tf.keras.models.load_model('Model.h5')

# Cargar el conjunto de datos para ajuste del preprocesador
file_path = 'heart.csv'
df = pd.read_csv(file_path)

# Identificar columnas numéricas y categóricas
num_features = ['Age', 'RestingBP', 'Cholesterol', 'MaxHR', 'Oldpeak']
cat_features = ['Sex', 'ChestPainType', 'FastingBS', 'RestingECG', 'ExerciseAngina', 'ST_Slope']

# Preprocesamiento: Escalar características numéricas y codificar características categóricas
preprocessor = ColumnTransformer(
    transformers=[
        ('num', StandardScaler(), num_features),
        ('cat', OneHotEncoder(handle_unknown='ignore'), cat_features)
    ])

# Ajustar el preprocesador al conjunto de datos
X = df.drop('HeartDisease', axis=1)
preprocessor.fit(X)

class HeartDiseaseInput(BaseModel):
    age: float
    sex: str
    chest_pain_type: str
    resting_bp: float
    cholesterol: float
    fasting_bs: int
    resting_ecg: str
    max_hr: float
    exercise_angina: str
    oldpeak: float
    st_slope: str

@app.get("/")
def read_root():
    return {"Hello": "World!"}

@app.post("/predict")
def predict(input: HeartDiseaseInput):
    try:
        # Crear un DataFrame con los datos de entrada
        input_data = pd.DataFrame({
            'Age': [input.age],
            'Sex': [input.sex],
            'ChestPainType': [input.chest_pain_type],
            'RestingBP': [input.resting_bp],
            'Cholesterol': [input.cholesterol],
            'FastingBS': [input.fasting_bs],
            'RestingECG': [input.resting_ecg],
            'MaxHR': [input.max_hr],
            'ExerciseAngina': [input.exercise_angina],
            'Oldpeak': [input.oldpeak],
            'ST_Slope': [input.st_slope]
        })

        # Aplicar preprocesamiento
        input_data_preprocessed = preprocessor.transform(input_data)

        # Realizar la predicción
        prediction_prob = model.predict(input_data_preprocessed)
        predicted_class = (prediction_prob > 0.3).astype(int)

        # Interpretar la predicción
        result = "Tiene enfermedad cardíaca." if predicted_class == 1 else "No tiene enfermedad cardíaca."

        return {"prediction": result, "probability": prediction_prob[0][0]}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))