import gradio as gr
import re
import emoji
import spacy
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from collections import Counter

# --- 1. CONFIGURACIÓN ---
MODEL_ID = "c-armor/finetuned_emotions"
DEVICE = "cpu" # Usamos CPU
EMOTION_LABELS = ["alegria", "tristeza", "ira", "miedo", "sorpresa", "asco"]
UMBRALES = {
    "alegria": 0.50, 
    "tristeza": 0.50, 
    "ira": 0.50,
    "miedo": 0.30, 
    "sorpresa": 0.30, 
    "asco": 0.30   # <--- Coincide con entrenamiento
}

# --- 2. CARGAR MODELOS (Se cargan UNA VEZ al iniciar el Space) ---
print("Cargando modelo spaCy...")
NLP = spacy.load("es_core_news_md")
STOP_WORDS = NLP.Defaults.stop_words
print("spaCy cargado y STOP_WORDS listas")

print(f"Cargando modelo BETO desde {MODEL_ID}...")
TOKENIZER = AutoTokenizer.from_pretrained("dccuchile/bert-base-spanish-wwm-cased", use_fast=False)
MODEL = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
MODEL.to(DEVICE)
MODEL.eval()
print(f"BETO cargado exitosamente en {DEVICE}")

# --- 3. FUNCIONES DE PREPROCESAMIENTO---
def preprocesar_texto(texto: str) -> str:
    if not isinstance(texto, str): return ""
    texto = texto.lower()
    texto = emoji.demojize(texto, language="es")
    texto = re.sub(r"https://?://\S+|www\.\S+", "", texto)
    texto = re.sub(r"[^a-záéíóúüñ_.,!?¿¡]", " ", texto)
    texto = re.sub(r"\s{2,}", " ", texto).strip()
    return texto

def lematizar_texto_para_lista(texto: str):
    """Lematiza el texto y devuelve una LISTA (para nube de palabras)"""
    with NLP.disable_pipes("parser", "ner"):
        doc = NLP(texto)
    return [token.lemma_ for token in doc if token.text.strip()]

# --- 4. FUNCIÓN API 1: ANÁLISIS DE EMOCIONES ---
def analizar_emociones(texto: str):
    if not texto or not isinstance(texto, str) or len(texto.strip()) < 3:
        return {"error": "texto_invalido"}
    try:
        texto_limpio = preprocesar_texto(texto)
        # Pasamos directo el texto_limpio al tokenizer
        inputs = TOKENIZER(texto_limpio, padding=True, truncation=True, max_length=128, return_tensors="pt").to(DEVICE)
        
        with torch.no_grad():
            outputs = MODEL(**inputs)
            logits = outputs.logits
        
        probabilidades = torch.sigmoid(logits).cpu().numpy()[0]
        emociones_detectadas = []
        probabilidades_dict = {}
        
        for i, emocion in enumerate(EMOTION_LABELS):
            prob = float(probabilidades[i])
            umbral = UMBRALES[emocion]
            probabilidades_dict[emocion] = round(prob, 4)
            if prob >= umbral:
                emociones_detectadas.append(emocion)
        
        if not emociones_detectadas:
            emociones_detectadas = ["neutro"]
        
        return {"emociones": emociones_detectadas, "probabilidades": probabilidades_dict}
    except Exception as e:
        print(f"Error en predicción de emociones: {e}")
        return {"error": "error_prediccion"}

# --- 5. FUNCIÓN API 2: FRECUENCIA DE PALABRAS ---
def generar_frecuencia_palabras(texto: str):
    """Calcula las 50 palabras más frecuentes (devuelve una lista de listas)"""
    if not texto or not isinstance(texto, str) or not texto.strip():
        return []
    try:
        lemmas = lematizar_texto_para_lista(preprocesar_texto(texto))
        # Usa las STOP_WORDS globales
        important_words = [lem for lem in lemmas if lem not in STOP_WORDS and lem.isalpha() and len(lem) > 2]
        word_counts = Counter(important_words)
        # Devuelve el formato que tu frontend espera (lista de [palabra, conteo])
        return word_counts.most_common(50)
    except Exception as e:
        print(f"Error en frecuencia de palabras: {e}")
        return []


# --- 6. LANZAR AMBAS APIS USANDO GR.BLOCKS ---
with gr.Blocks() as demo:
    # --- Interfaz/API para Emociones ---
    with gr.Tab("Análisis de Emociones"):
        gr.Interface(
            fn=analizar_emociones,
            inputs=gr.Textbox(lines=5, label="Texto a Analizar"),
            outputs="json",
            api_name="predict_emotions" # <-- API endpoint 1
        )
    
    # --- Interfaz/API para Frecuencia de Palabras ---
    with gr.Tab("Frecuencia de Palabras"):
        gr.Interface(
            fn=generar_frecuencia_palabras,
            inputs=gr.Textbox(lines=5, label="Texto completo"),
            outputs="json",
            api_name="predict_frequency" # <-- API endpoint 2
        )

demo.launch()