import streamlit as st
from transformers import pipeline
import matplotlib.pyplot as plt
import io
import json

# Load the emotion classification model
@st.cache_resource
def load_model():
    return pipeline(
        "text-classification", 
        model="j-hartmann/emotion-english-distilroberta-base", 
        return_all_scores=True
    )

emotion_classifier = load_model()

# Define a function to predict emotions and generate a bar chart
def predict_emotion_with_chart(text):
    if not text.strip():
        return None, None

    # Get predictions
    results = emotion_classifier(text)
    emotions = {result['label']: round(result['score'], 2) for result in results[0]}
    sorted_emotions = dict(sorted(emotions.items(), key=lambda item: item[1], reverse=True))

    # Create a bar chart
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(sorted_emotions.keys(), sorted_emotions.values(), color="skyblue")
    ax.set_title("Emotion Scores")
    ax.set_ylabel("Confidence Score")
    ax.set_ylim(0, 1)
    ax.set_xticklabels(sorted_emotions.keys(), rotation=45, ha="right")
    plt.tight_layout()

    return sorted_emotions, fig

# Function to generate JSON result
def generate_json(text):
    results = emotion_classifier(text)
    emotions = {result['label']: round(result['score'], 2) for result in results[0]}
    return json.dumps(emotions, indent=2)

# Streamlit UI
st.title("🌟 Enhanced Emotion Detection App")
st.markdown("Analyze the emotions in a sentence and visualize them. Enter text to detect emotions, see a bar chart of scores, and download the results as a JSON file.")

# Input text
text_input = st.text_area("Enter text to analyze emotions:", "")

if st.button("Analyze Emotions"):
    if text_input.strip():
        emotion_scores, chart = predict_emotion_with_chart(text_input)
        if emotion_scores:
            st.subheader("Emotion Scores")
            st.json(emotion_scores)
            st.subheader("Emotion Chart")
            st.pyplot(chart)
        else:
            st.error("No emotions detected. Please enter a valid text.")
    else:
        st.warning("Please enter some text.")

# Download JSON results
if text_input.strip():
    json_data = generate_json(text_input)
    st.download_button(label="Download Results as JSON", data=json_data, file_name="emotion_results.json", mime="application/json")