from difflib import Differ

import gradio as gr
from transformers import pipeline

pipe = pipeline("summarization", "dominguesm/positive-reframing-ptbr")


def predict(text, operation):
    try:
        res = pipe(f"[{operation}]: {text}", max_length=128)
    except Exception as e:
        return e

    d = Differ()
    return (
        res[0]["summary_text"],
        [
            (token[2:], token[0] if token[0] != " " else None)
            for token in d.compare(text, res[0]["summary_text"])
        ],
    )
    # return res[0]["summary_text"]


iface = gr.Interface(
    title="Positive Reframing PT-BR",
    description="This model is a PTT5 adjusted to the sentiment transfer task, where the objective is to reverse the sentiment polarity of a text without contradicting the original meaning. Positive reframing induces a complementary positive viewpoint (e.g. glass-half-full) escaping negative patterns. More info [here](https://huggingface.co/dominguesm/positive-reframing-ptbr).",
    fn=predict,
    inputs=[
        gr.Textbox(
            lines=1,
            placeholder=(
                f"Pensar no meu futuro me faz querer viver numa ilha sozinha para sempre"
            ),
        ),
        gr.Radio(
            [
                "growth",
                "impermanence",
                "neutralizing",
                "optimism",
                "self_affirmation",
                "thankfulness",
            ]
        ),
    ],
    outputs=[
        gr.Textbox(label="Generated Text"),
        gr.HighlightedText(
            label="Diff",
            combine_adjacent=True,
        ).style(color_map={"+": "green", "-": "red"}),
    ],
    examples=[
        [
            "Tenho tanta coisa para fazer antes de sair da cidade por uma semana no domingo.",
            "thankfulness",
        ],
        [
            "Aquele momento, você percebe que todo o trabalho duro foi para nada e você não tem controle sobre nada.",
            "self_affirmation",
        ],
        [
            "Um daqueles dias em que você sente que precisa de um abraço de cinco minutos, um pacote de marshmallows e um amigo para conversar.",
            "impermanence",
        ],
    ],
)

iface.launch()