from fastapi import FastAPI, Query
import joblib
import pandas as pd

# Cargar modelo y datos
modelo_probabilidad = joblib.load("modelo_probabilidad_paro.pkl")
features_probabilidad = joblib.load("features_modelo_paro.pkl")
df_probabilidad = joblib.load("df_codificado_paro.pkl")

app = FastAPI()

@app.get("/")
def root():
    return {"mensaje": "API de predicción de probabilidad"}

@app.get("/probabilidad")
def probabilidad(maquina: str = Query(..., description="Nombre de la máquina")):
    if maquina not in df_probabilidad["Maquina"].unique():
        return {"error": f"Máquina '{maquina}' no encontrada en datos de probabilidad."}
    try:
        maquina_encoded = df_probabilidad[df_probabilidad["Maquina"] == maquina]["Maquina_encoded"].iloc[0]
        df_maquina = df_probabilidad[df_probabilidad["Maquina"] == maquina]

        fila_test = df_maquina[features_probabilidad].mean().to_frame().T
        fila_test["Maquina_encoded"] = maquina_encoded

        probas = modelo_probabilidad.predict_proba(fila_test)[0]
        prob_marcha = probas[1]
        prob_paro = probas[0]

        return {
            "tipo": "probabilidad",
            "maquina": maquina,
            "probabilidad_paro": round(prob_paro * 100, 2),
            "probabilidad_marcha": round(prob_marcha * 100, 2)
        }
    except Exception as e:
        return {"error": f"Error en predicción de probabilidad: {str(e)}"}