import streamlit as st

from pipeline.model import MultiTaskModel
from pipeline.preprocessing import Preprocessor

# Load the model
EMOTION_CHOICES = (
    "Angry",
    "Disgust",
    "Happy",
    "Neutral",
    "Sad",
    "Surprise",
)
TOXICITY_CHOICES = (
    "Blaming Others",
    "Cyberbullying",
    "Gameplay Experience Complaints",
    "Gamesplaining",
    "Multiple Discrimination",
    "Not Toxic",
    "Sarcasm",
)

st.title("Emotion and Toxicity Classification of Valorant chat messages")

st.write(
    'This is a simple web app that predicts the emotion and toxicity of Valorant chat messages. Enter a message in the text box below and click the "Predict" button to get the prediction.'
)

st.table(
    {
        "Emotion": EMOTION_CHOICES,
        "Toxicity": TOXICITY_CHOICES,
    }
)

@st.cache_resource
def loading_model():
    return MultiTaskModel(preprocessor=Preprocessor())

model = loading_model()

# Get user input
user_input = st.text_input("Enter a Valorant chat message:")
st.write("You entered:", user_input)

# Predict
prediction = model.predict(user_input)
emotions, toxicitys = prediction

col1, col2 = st.columns(2)

with col1:
    for i, emotion in enumerate(emotions[0]):
        st.write(f"{EMOTION_CHOICES[i]}: {(emotion*100):.2f}%")
        st.progress(float(emotion))
with col2:
    for i, toxicity in enumerate(toxicitys[0]):
        st.write(f"{TOXICITY_CHOICES[i]}: {(toxicity*100):.2f}%")
        st.progress(float(toxicity))

decoded = model.decode(prediction)
# Display the prediction
st.write("The predicted emotion is:", decoded[0][0])
st.write("The predicted toxicity is:", decoded[1][0])