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"""
    split_sentence = x.split(" ")
    for i_word, word in enumerate(split_sentence):
        yield word, " ".join(split_sentence[:i_word] + split_sentence[i_word + 1:])


def racism_analysis(text):
    results = racism_analysis_pipe(text, return_all_scores=True)[0][1]
    label = "Non-racist" if results["score"] < 0.5 else "Racist"
    score = results["score"]

    words = {}
    colors = []
    for word, sentence in iterate_sentences(text):
        words[word] = score / \
            racism_analysis_pipe(sentence, return_all_scores=True)[
            0][1]["score"]
        if label == "Non-racist":
            color = None
        elif words[word] > 10:
            color = "Very racist"
        elif words[word] > 2:
            color = "Racist"
        else:
            color = 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(
            label="Racist heatmap", color_map={"Very racist": "red", "Racist": "pink"}),
    ],
    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"]
    ],
    flagging_options=["Correct label", "Incorrect label"],
    allow_flagging="manual",
)

gradio_ui.launch()