import gradio as gr
from huggingface_hub import hf_hub_download
from fastai.learner import load_learner

# 1. Descargar el modelo FastAI desde tu repositorio de modelos
REPO_ID = "Camayli/practica7"
FILENAME = "model.pkl"

print("Descargando el modelo...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

# 2. Cargar el modelo usando FastAI
print("Cargando el modelo FastAI...")
learn = load_learner(model_path)

# 3. Definir la función de predicción
def predict_emotion(text: str):
    text = (text or "").strip()
    if not text:
        return {"error": "Escribe un texto para clasificar."}, ""
    
    # FastAI devuelve la clase predicha, el índice y las probabilidades
    pred_class, pred_idx, probabilities = learn.predict(text)
    
    # Obtener la lista de etiquetas (emociones)
    # En FastAI para texto, el vocabulario suele ser una tupla donde la posición 1 tiene las clases
    if isinstance(learn.dls.vocab, tuple):
        classes = list(learn.dls.vocab[1])
    else:
        classes = list(learn.dls.vocab)
        
    # Convertir a un diccionario de {etiqueta: probabilidad} para Gradio
    scores = {str(classes[i]): float(probabilities[i]) for i in range(len(classes))}
    
    # Formatear el resumen de salida
    summary = f"Predicción: {pred_class} | confianza: {float(probabilities[pred_idx]):.4f}"
    
    return scores, summary

# 4. Construir la interfaz de Gradio
DESCRIPTION = """Este Space clasifica texto usando un modelo FastAI entrenado en el ejercicio obligatorio.
Pega una frase en inglés y el sistema devolverá la emoción predicha."""

EXAMPLES = [
    ["i feel very happy today"],
    ["i am extremely sad and disappointed"],
    ["i am angry about what happened"],
    ["i feel scared and nervous"],
    ["i feel so much love for my family"],
    ["i am surprised by the news"],
]

with gr.Blocks() as demo:
    gr.Markdown("# Clasificador de emociones")
    gr.Markdown(DESCRIPTION)
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Texto de entrada",
                placeholder="Write a sentence here...",
                lines=4,
            )
            predict_btn = gr.Button("Clasificar")
        with gr.Column(scale=2):
            label_output = gr.Label(label="Probabilidades")
            summary_output = gr.Textbox(label="Resultado", interactive=False)
            
    gr.Examples(examples=EXAMPLES, inputs=text_input)
    
    predict_btn.click(
        fn=predict_emotion,
        inputs=text_input,
        outputs=[label_output, summary_output],
    )
    text_input.submit(
        fn=predict_emotion,
        inputs=text_input,
        outputs=[label_output, summary_output],
    )

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