import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BertForSequenceClassification, BertTokenizer


# --- MODELO DE SENTIMIENTOS ---
sentiment_tokenizer = BertTokenizer.from_pretrained("VerificadoProfesional/SaBERT-Spanish-Sentiment-Analysis")
sentiment_model = BertForSequenceClassification.from_pretrained("VerificadoProfesional/SaBERT-Spanish-Sentiment-Analysis")
sentiment_pipeline = pipeline("sentiment-analysis", model=sentiment_model, tokenizer=sentiment_tokenizer)

# --- MODELO DE REESCRITURA GPT-2 ---
rewrite_tokenizer = AutoTokenizer.from_pretrained("flax-community/gpt-2-spanish")
rewrite_model = AutoModelForCausalLM.from_pretrained("flax-community/gpt-2-spanish")
rewrite_pipeline = pipeline("text-generation", model=rewrite_model, tokenizer=rewrite_tokenizer)

# --- FUNCIÓN PRINCIPAL ---
def analizar_y_reescribir(tweet, max_retries=10):
    sentimiento = sentiment_pipeline(tweet)[0]
    etiqueta = sentimiento["label"]
    score = sentimiento["score"]

    reescritura = ""
    if etiqueta == "Negative":
        for intento in range(max_retries):
            prompt = f"Este es el tweet: {tweet}\nEscribe una respuesta positiva y motivadora:\n"
            resultado = rewrite_pipeline(
                prompt,
                max_new_tokens=60,
                do_sample=True,
                temperature=0.7,
                top_k=50,
                top_p=0.9,
                pad_token_id=rewrite_tokenizer.eos_token_id,
                return_full_text=False
            )[0]["generated_text"]
            reescritura = resultado.strip().split("\n")[0]

            # Revisar sentimiento de la propuesta
            senti2 = sentiment_pipeline(reescritura)[0]
            if senti2["label"] == "Positive":
                break
            # Si no salió positivo, sigue al siguiente intento

    return f"{etiqueta} ({score:.2f})", reescritura

# --- INTERFAZ GRADIO ---
gr.Interface(
    fn=analizar_y_reescribir,
    inputs=gr.Textbox(label="Introduce un tweet", placeholder="Por ejemplo: Este servicio es terrible y no lo recomiendo."),
    outputs=[
        gr.Textbox(label="Sentimiento detectado en el tweet"),
        gr.Textbox(label="Frase positiva o motivadora")
    ],
    title="Análisis de Tweets Negativos",
).launch()