import gradio as gr
from transformers import (
    pipeline,
    AutoModelForSequenceClassification,
    AutoTokenizer
)
import torch
import numpy as np

# Configurar el dispositivo
device = "cuda" if torch.cuda.is_available() else "cpu"

# === CONFIGURACIÓN DEL CHATBOT ===
chat_generator = pipeline(
    'text-generation',
    model='microsoft/DialoGPT-small',
    device=device
)

# === CONFIGURACIÓN DEL ANALIZADOR DE SENTIMIENTOS ===
model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
sentiment_tokenizer = AutoTokenizer.from_pretrained(model_name)
sentiment_model = AutoModelForSequenceClassification.from_pretrained(model_name)
sentiment_model.to(device)

def chatbot(mensaje):
    if not mensaje.strip():
        return "Por favor, escribe un mensaje.", None, None

    try:
        # Generar respuesta del chatbot
        respuesta = chat_generator(
            mensaje,
            max_length=100,
            temperature=0.7,
            do_sample=True,
            top_p=0.95,
            num_return_sequences=1
        )[0]['generated_text']

        respuesta = respuesta.replace(mensaje, "").strip()

        # Analizar sentimiento del mensaje del usuario
        inputs = sentiment_tokenizer(mensaje, return_tensors="pt", truncation=True,
                                   max_length=512, padding=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        outputs = sentiment_model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        rating = torch.argmax(predictions).item() + 1
        confidence = predictions[0][rating-1].item()

        if rating <= 2:
            sentimiento = "Muy Negativo"
        elif rating == 3:
            sentimiento = "Neutral"
        elif rating == 4:
            sentimiento = "Positivo"
        else:
            sentimiento = "Muy Positivo"

        sentimiento_completo = f"{sentimiento} ({rating} estrellas)"
        confianza = round(confidence * 100, 2)

        return respuesta, sentimiento_completo, confianza

    except Exception as e:
        return f"Error: {str(e)}", "Error en el análisis", 0.0

# Crear la interfaz
demo = gr.Interface(
    fn=chatbot,
    inputs=[
        gr.Textbox(
            placeholder="Escribe tu mensaje...",
            label="Mensaje",
            lines=3
        )
    ],
    outputs=[
        gr.Textbox(label="Respuesta del Chatbot"),
        gr.Label(label="Sentimiento de tu mensaje"),
        gr.Number(label="Confianza del análisis (%)")
    ],
    title="Chatbot con Análisis de Sentimientos",
    description="Un chatbot que responde a tus mensajes y analiza el sentimiento de lo que escribes",
    examples=[
        ["¡Estoy muy feliz hoy!"],
        ["No me gusta nada este servicio, es terrible."],
        ["El día está normal, nada especial que contar."],
        ["¡Me encanta hablar contigo!"],
        ["Estoy un poco decepcionado con los resultados."]
    ],
    allow_flagging="never",
    cache_examples=True
)

# Lanzar la interfaz
demo.launch(share=True, debug=True)