import streamlit as st
import transformers
import numpy as np
import pandas as pd
from scipy.special import softmax
import torch


def preprocess(text):
    new_text = []
    for t in text.split(" "):
        t = '@user' if t.startswith('@') and len(t) > 1 else t
        t = 'http' if t.startswith('http') else t
        new_text.append(t)
    return " ".join(new_text)


st.title("Toxicity Classification App")

user_input = st.text_input("Input texts to analyze", "Cody Jiang is a happy boy!")

model_names = ['distilbert-base-uncased-finetuned-sst-2-english', 'Codys-Finetuning-Language-Model']
model_name = st.selectbox("Select a pretrained model", model_names)

if model_name == "Codys-Finetuning-Language-Model":
    model = transformers.BertForSequenceClassification.from_pretrained("./myModel/")
    tokenizer = transformers.BertTokenizerFast.from_pretrained("bert-base-uncased")
else:
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
    config = transformers.AutoConfig.from_pretrained(model_name)
    model = transformers.AutoModelForSequenceClassification.from_pretrained(model_name)

if st.button("Submit"):
    if model_name == "Codys-Finetuning-Language-Model":
        input_val = tokenizer(user_input, padding=True, truncation=True, max_length=512, return_tensors="pt")
        output_val = model(**input_val)
        probabilities = torch.sigmoid(output_val.logits)
        result_list = probabilities.tolist()[0]
        columns = ["Tweet", "Toxicity Class", "Probability"]
        toxicity_class = ["toxic", "severe toxic", "obscene", "threat", "insult", "identity hate"]
        result_df = pd.DataFrame(columns=columns)
        result_df.loc[0] = [user_input, toxicity_class[result_list.index(max(result_list))], max(result_list)]
        st.table(result_df)

    else:
        text = preprocess(user_input)
        encoded_input = tokenizer(text, return_tensors='pt')
        output = model(**encoded_input)
        scores = output[0][0].detach().numpy()
        scores = softmax(scores)

        ranking = np.argsort(scores)
        ranking = ranking[::-1]
        columns = ["Tweet", "Toxicity Class", "Probability"]
        toxicity_class = config.id2label
        result_list = []
        columns_list = []
        for i in range(scores.shape[0]):
            l = toxicity_class[ranking[i]]
            s = scores[ranking[i]]
            result_list.append(s)
            columns_list.append(l)
        result_df = pd.DataFrame(columns=columns)
        result_df.loc[0] = [user_input, columns_list[0], result_list[0]]
        st.table(result_df)