import gradio as gr
import joblib
import pandas as pd

# Cargar modelos entrenados
MODELOS = {
    "Femenino - Random Forest": joblib.load("fem_modelos_rf_sin_leakage.pkl"),
    "Femenino - XGBoost": joblib.load("fem_modelos_xgb_sin_leakage.pkl"),
    "FIBA - Random Forest": joblib.load("fiba_modelos_rf_sin_leakage.pkl"),
    "FIBA - XGBoost": joblib.load("fiba_modelos_xgb_sin_leakage.pkl"),
    "NBA - Random Forest": joblib.load("nba_modelos_rf_sin_leakage.pkl"),
    "NBA - XGBoost": joblib.load("nba_modelos_xgb_sin_leakage.pkl"),
}

# MAEs históricos (ajusta si cambian)
MAES = {
    "Femenino - Random Forest": {"3er cuarto": 4.50, "4to cuarto": 4.15, "Tiempo completo": 4.49},
    "Femenino - XGBoost": {"3er cuarto": 4.66, "4to cuarto": 4.65, "Tiempo completo": 4.86},
    "FIBA - Random Forest": {"3er cuarto": 4.46, "4to cuarto": 4.38, "Tiempo completo": 4.52},
    "FIBA - XGBoost": {"3er cuarto": 4.63, "4to cuarto": 4.67, "Tiempo completo": 4.85},
    "NBA - Random Forest": {"3er cuarto": 5.35, "4to cuarto": 4.88, "Tiempo completo": 5.06},
    "NBA - XGBoost": {"3er cuarto": 5.57, "4to cuarto": 5.13, "Tiempo completo": 5.36},
}

def maes_table():
    maes_df = pd.DataFrame(MAES).T
    maes_df.index.name = "Modelo"
    return maes_df.reset_index()

def predecir_puntos(
    modelo_nombre,
    home_1, home_2, home_3,
    away_1, away_2, away_3
):
    h1 = home_1 if home_1 is not None else 0
    h2 = home_2 if home_2 is not None else 0
    h3 = home_3 if home_3 is not None else 0
    a1 = away_1 if away_1 is not None else 0
    a2 = away_2 if away_2 is not None else 0
    a3 = away_3 if away_3 is not None else 0

    modelos = MODELOS[modelo_nombre]

    pred_3 = pred_4 = pred_total = ""
    # Predicción del 3er cuarto (si solo hay 1er y 2do cuarto)
    if (h3 == 0 and a3 == 0) and (h1 != 0 or h2 != 0 or a1 != 0 or a2 != 0):
        pace = h1 + h2 + a1 + a2
        X3 = [[h1, h2, a1, a2, pace]]
        pred_3 = modelos["tercer_cuarto"].predict(X3)[0]
        pred_3 = f"{pred_3:.2f} puntos (predicción 3er cuarto)"
    else:
        pred_3 = "No disponible (ya jugado o faltan datos)"

    # Predicción del 4to cuarto y total (si hay hasta 3er cuarto)
    if h3 != 0 or a3 != 0:
        pace = h1 + h2 + h3 + a1 + a2 + a3
        X4 = [[h1, h2, h3, a1, a2, a3, pace]]
        pred_4 = modelos["cuarto_cuarto"].predict(X4)[0]
        pred_4 = f"{pred_4:.2f} puntos (predicción 4to cuarto)"
        pred_total = modelos["tiempo_completo"].predict(X4)[0]
        pred_total = f"{pred_total:.2f} puntos (predicción tiempo completo)"
    else:
        pred_4 = "No disponible (faltan datos de 3er cuarto)"
        pred_total = "No disponible (faltan datos de 3er cuarto)"

    return pred_3, pred_4, pred_total, maes_table()

with gr.Blocks() as demo:
    gr.Markdown("# Predicción de Puntos por Cuarto y Tiempo Completo")
    gr.Markdown(
        "Selecciona la liga/modelo, ingresa los puntos de los cuartos jugados para ambos equipos (deja en blanco los que no han ocurrido)."
    )
    with gr.Row():
        modelo_input = gr.Dropdown(choices=list(MODELOS.keys()), label="Liga y Modelo")
    with gr.Row():
        home_1_input = gr.Number(label="Puntos Home 1er cuarto", value=None)
        home_2_input = gr.Number(label="Puntos Home 2do cuarto", value=None)
        home_3_input = gr.Number(label="Puntos Home 3er cuarto", value=None)
    with gr.Row():
        away_1_input = gr.Number(label="Puntos Away 1er cuarto", value=None)
        away_2_input = gr.Number(label="Puntos Away 2do cuarto", value=None)
        away_3_input = gr.Number(label="Puntos Away 3er cuarto", value=None)
    with gr.Row():
        pred3_output = gr.Textbox(label="Predicción 3er cuarto")
        pred4_output = gr.Textbox(label="Predicción 4to cuarto")
        predtotal_output = gr.Textbox(label="Predicción tiempo completo")
    gr.Markdown("## Tabla de MAEs históricos")
    maes_output = gr.Dataframe(value=maes_table(), interactive=False)

    def wrapper(modelo, h1, h2, h3, a1, a2, a3):
        p3, p4, pt, maes_df = predecir_puntos(modelo, h1, h2, h3, a1, a2, a3)
        return p3, p4, pt, maes_df

    gr.Button("Predecir").click(
        wrapper,
        inputs=[modelo_input, home_1_input, home_2_input, home_3_input, away_1_input, away_2_input, away_3_input],
        outputs=[pred3_output, pred4_output, predtotal_output, maes_output]
    )

if __name__ == "__main__":
    demo.launch()