import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from collections import Counter
from scipy.special import softmax

article_string = "Author: <a href=\"https://huggingface.co/FpOliveira\">Felipe Ramos de Oliveira</a>. Read more about our <a href=\"https://github.com/Silly-Machine/TuPi-Portuguese-Hate-Speech-Dataset\">The Portuguese hate speech dataset (TuPI) </a>."

app_title = "Portuguese hate speech classifier - Classicador de discurso de ódio em português"

app_description = """
EN: This application employs multiple models to identify hate speech in Portuguese texts. You have the option to enter your own phrases by filling in the "Text" field or choosing one of the examples provided below.
\nPT: Esta aplicativo emprega múltiplos modelos para identificar discurso de ódio em textos portugueses. Você tem a opção de inserir suas próprias frases preenchendo o campo "Texto" ou escolhendo um dos exemplos abaixo
"""

app_examples = [
    ["bom dia flor do dia!!!"],
    ["o ódio é muito grande no coração da ex-deputada federal joise hasselmann contra a família bolsonaro"],
    ["mano deus me livre q nojo da porra!🤮🤮🤮🤮🤮"],
    ["obrigada princesa, porra, tô muito feliz snrsss 🤩🤩🤩❤️"],
    ["mds mas o viado vir responder meus status falando q a taylor foi racista foi o auge 😂😂"],
    ["Pra ser minha inimiga no mínimo tem que ter um rostinho bonito e delicado, não se considere minha rival com essa sua cara de cavalo não, feia, cara de traveco, cabeçuda, queixo quadrado 🤣🤣"]
]

output_textbox_component_description = """
EN: This box will display hate speech results based on the average score of multiple models.
PT: Esta caixa exibirá resultados da classicação de discurso de ódio com base na pontuação média de vários modelos.
"""

output_json_component_description = { "breakdown": """
This box presents a detailed breakdown of the evaluation for each model.
""",
"detalhamento": """
(Esta caixa apresenta um detalhamento da avaliação para cada modelo.)
""" }

short_score_descriptions = {
   0: "Not hate",
   1: "Hate"
}

# Define hate speech categories
hate_speech_categories = {
    0: "ageism",
    1: "aporophobia",
    2: "body_shame",
    3: "capacitism",
    4: "lgbtphobia",
    5: "political",
    6: "racism",
    7: "religious intolerance",
    8: "misogyny",
    9: "xenophobia",
    10: "other",
    11: "not hate"
}
model_list = [
    "FpOliveira/tupi-bert-large-portuguese-cased",
    "FpOliveira/tupi-bert-base-portuguese-cased",
]

user_friendly_name = {
    "FpOliveira/tupi-bert-large-portuguese-cased": "BERTimbau large (TuPi)",
    "FpOliveira/tupi-bert-base-portuguese-cased":  "BERTimbau base (TuPi)",
}

reverse_user_friendly_name = { v:k for k,v in user_friendly_name.items() }

user_friendly_name_list = list(user_friendly_name.values())

model_array = []

for model_name in model_list:
    row = {}
    row["name"] = model_name
    row["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
    row["model"] = AutoModelForSequenceClassification.from_pretrained(model_name)
    model_array.append(row)
 
def most_frequent(array):
    occurence_count = Counter(array)
    return occurence_count.most_common(1)[0][0]


def predict(s1, chosen_model):
    if not chosen_model:
        chosen_model = user_friendly_name_list[0]
    full_chosen_model_name = reverse_user_friendly_name[chosen_model]
    
    for row in model_array:
        name = row["name"]
        if name != full_chosen_model_name:
            continue
        else:
            tokenizer = row["tokenizer"]
            model = row["model"]
            model_input = tokenizer(*([s1],), padding=True, return_tensors="pt")
            with torch.no_grad():
                output = model(**model_input)
                logits = output[0][0].detach().numpy()
                logits = softmax(logits).tolist()
                break

    # Get the indices of the top two probabilities
    top_two_indices = sorted(range(len(logits)), key=lambda i: logits[i], reverse=True)[:2]

    # Get the categories and probabilities for the top two
    top_two_categories = [hate_speech_categories[str(index)] for index in top_two_indices]
    top_two_probabilities = [logits[index] for index in top_two_indices]

    result = {
        "predicted_categories": top_two_categories,
        "probabilities": top_two_probabilities,
    }

    return result

inputs = [
    gr.Textbox(label="Text", value=app_examples[0][0]),
    gr.Dropdown(label="Model", choices=user_friendly_name_list, value=user_friendly_name_list[0])
]

# Output components
outputs = [
    gr.Label(label="Top Predicted Categories"),
    gr.Label(label="Top Probabilities"),
]

gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=app_title, 
             description=app_description,
             examples=app_examples,
             article = article_string).launch()