import os

import gradio as gr
from transformers import pipeline

RACISM_MODEL = "davidmasip/racism"
racism_analysis_pipe = pipeline(
    "text-classification",
    model=RACISM_MODEL,
    tokenizer=RACISM_MODEL,
    use_auth_token=os.getenv("access_token")
)


def iterate_sentences(x):
    """For each word of the sentence, returns the sentence without the word"""
    for word in x.split(" "):
        yield word, x.replace(word, "")


def racism_analysis(text):
    results = racism_analysis_pipe(text)[0]
    label = "Non-racist" if results["label"] == "LABEL_0" else "Racist"
    score = (
        1 - results["score"]
        if results["label"] == "LABEL_0"
        else results["score"]
    )

    words = {}
    colors = []
    for word, sentence in iterate_sentences(text):
        words[word] = 1 / \
            racism_analysis_pipe(sentence, return_all_scores=True)[
            0][1]["score"]
        color = "+" if words[word] > 100 else None
        colors.append((word, color))

    print(colors)
    print(words)
    return label, round(100 * score), colors


gradio_ui = gr.Interface(
    fn=racism_analysis,
    title="Racism Detector (Spanish)",
    description="Enter some text and check if model detects racism.",
    inputs=[
        gr.inputs.Textbox(lines=5, label="Paste some text here"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Label"),
        gr.outputs.Textbox(label="Racism score (0 - 100)"),
        gr.outputs.HighlightedText(color_map={"+": "red"}),
    ],
    examples=[
        ["Unos menas roban a una mujer"],
        ["Unos chinos roban a una mujer"],
        ["Unos árabes roban a una mujer"],
        ["Unos moros roban a una mujer"],
        ["Unos latinos roban a una mujer"],
        ["No me gustan los menas"],
        ["No me gustan los chinos"],
        ["No me gustan los árabes"],
        ["No me gustan los moros"],
        ["No me gustan los latinos"],
        ["El gobierno levanta el confinamiento"]
    ],
)

gradio_ui.launch()