import streamlit as st
from stereotype_detector import Detector

st.title("Stereotype Detection")

EXAMPLES = {
    "Gender": "The baby loved the presence of his caring mommy.",
    "Religion": "The orange church had a very good reputation.",
    "Race": "Steve thinks singaporeans are short, but he doesn't know any better",
    "Profession": "The football player's draft card said he was tough.",
    "Input your own": ""
}
if "detector" not in st.session_state:
    st.session_state["detector"] = None

def format_results(results, bias_level):
    formatted = ""
    for result in results:
        for text, pred in result.items():
            formatted += f"**Text**: {text}\n\n"
            formatted += "**Predictions**:\n"
            if isinstance(pred, dict):
                for token, labels in pred.items():
                    if isinstance(labels, dict):
                        formatted += f"- Token: `{token}`\n"
                        for label, score in labels.items():
                            formatted += f"    - Label: `{label}`, Score: `{score}`\n"
                if bias_level == "Sentence":  # sort the labels only for sentence level bias
                    sorted_pred = dict(sorted(pred.items(), key=lambda item: item[1], reverse=True))
                    for label, score in sorted_pred.items():
                        formatted += f"- Label: `{label}`, Score: `{score}`\n"
            else:
                formatted += f"Prediction score: {pred}\n"
    return formatted

level = st.selectbox("Select the Detection Levels:", ("Sentence","Token"))

if st.button("Load Models"):
    with st.spinner('Loading models...'):
        st.session_state["detector"] = Detector(level)
        dummy_sentence = "This is a dummy sentence."
        dummy_result = st.session_state["detector"].predict([dummy_sentence])
        if dummy_result:
            st.success("Models loaded successfully!")
        else:
            st.error("Failed to load models. Please check the server and/or model parameters.")

example_type = st.selectbox("Choose an example type:", list(EXAMPLES.keys()))
target_sentence = st.text_input("Input the sentence you want to detect:", value=EXAMPLES[example_type])

if st.button("Detect"):
    with st.spinner('Detecting...'):
        results = st.session_state["detector"].predict([target_sentence])
        if results:
            formatted_results = format_results(results, level)  # pass the selected bias level to the function
            st.markdown(f"## Detection Results: \n\n {formatted_results}")
        else:
            st.error("Prediction failed. Please check the input and try again.")