import streamlit as st
from transformers import pipeline, AutoConfig

MODEL_NAME = "blaikhole/distilbert-review-bug-classifier"

label_mapping = {
    "LABEL_0": "Graphics issue 🎨",
    "LABEL_1": "Network issue 🌐",
    "LABEL_2": "No Bug ✅",
    "LABEL_3": "Performance issue 🚀"
}

# Load model config to get label mapping
config = AutoConfig.from_pretrained(MODEL_NAME)
id2label = config.id2label

# Create a pipeline for text classification
pipe = pipeline("text-classification", model=MODEL_NAME)

# Streamlit app UI
st.title("Review Bug Classification Demo 🐞")
st.write("Enter some text and the model will predict the bug category.")

# User Input
user_input = st.text_area("Input Text:", height=150)

# Prediction
if st.button("Classify"):
    if user_input:
        result = pipe(user_input, return_all_scores=True)[0]  # Get all scores

        predictions = {label_mapping.get(res['label'], res['label']): int(res['score'] * 100) for res in result}
        
        ordered_labels = ["Graphics issue 🎨", "Network issue 🌐", "No Bug ✅", "Performance issue 🚀"]
        ordered_predictions = {k: predictions[k] for k in ordered_labels if k in predictions}

        # Get top prediction
        top_label = max(ordered_predictions, key=ordered_predictions.get)

        # Show top category
        st.write(f"### 🏆 Predicted Category: `{top_label}`")

        st.write("### Confidence Scores:")
        for label, score in ordered_predictions.items():
            st.write(f"**{label}**")
            st.progress(score)  # Display confidence as a progress bar

    else:
        st.warning("⚠️ Please enter some text.")