import gradio as gr
import torch
import torch.nn.functional as F
import re

from transformers import BertTokenizer, BertForSequenceClassification

# Cargar modelo
model_path = "."

tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path)

# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# Aspectos
aspect_keywords = {
    "quality": ["quality", "design"],
    "price": ["price", "cheap", "expensive", "worth"],
    "shipping": ["shipping", "delivery", "arrive", "arrival", "took"]
}

# Predicción
def predecir_sentimiento(texto):

    inputs = tokenizer(
        texto,
        return_tensors="pt",
        truncation=True,
        padding=True
    )

    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits

    temperature = 2.5

    probs = F.softmax(logits / temperature, dim=1)

    pred = torch.argmax(logits, dim=1).item()

    confianza = probs.max().item()

    sentimiento = "positivo" if pred == 1 else "negativo"

    return sentimiento, confianza


# Aspectos
def analizar_aspectos(texto):

    frases = re.split(
        r'[.,;!]| but | and | however | although | though | because ',
        texto
    )

    resultado = {}

    for frase in frases:

        for asp, palabras in aspect_keywords.items():

            if any(p in frase for p in palabras):

                sentimiento, confianza = predecir_sentimiento(frase)

                resultado[asp] = {
                    "sentimiento": sentimiento,
                    "confianza": round(confianza, 2)
                }

    return resultado


# Función principal
def analizar_resena(texto):

    sentimiento, confianza = predecir_sentimiento(texto)
    aspectos = analizar_aspectos(texto)

    resultado = f"Sentimiento general: {sentimiento} ({confianza:.2f})\n\n"
    resultado += "Aspectos:\n"

    for asp, info in aspectos.items():
        resultado += f"- {asp}: {info['sentimiento']} ({info['confianza']})\n"

    return resultado


# Interfaz
app = gr.Interface(
    fn=analizar_resena,
    inputs=gr.Textbox(
        lines=6,
        placeholder="Write a review here...",
        label="Review"
    ),
    outputs=gr.Textbox(),
    title="Sentiment Analysis",
    description="Analyze product/services reviews using BERT"
)

app.launch()