import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import plotly.graph_objects as go

# Page config
st.set_page_config(
    page_title="Emotion Detector",
    page_icon="📊",
    layout="wide"
)

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
    model = AutoModelForSequenceClassification.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
    return tokenizer, model

def analyze_text(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    return probs[0].detach().numpy()

def create_emotion_plot(emotions_dict):
    fig = go.Figure(data=[
        go.Bar(
            x=list(emotions_dict.keys()),
            y=list(emotions_dict.values()),
            marker_color=['#FF9999', '#99FF99', '#9999FF', '#FFFF99', '#FF99FF', '#99FFFF', '#FFB366']
        )
    ])
    
    fig.update_layout(
        title="Emotion Analysis Results",
        xaxis_title="Emotions",
        yaxis_title="Confidence Score",
        yaxis_range=[0, 1]
    )
    return fig

# App title and description
st.title("📊 Text Emotion Analysis")
st.markdown("""
This app analyzes the emotional content of your text using a pre-trained emotion detection model.
Try typing or pasting some text below!
""")

# Load model
with st.spinner("Loading model..."):
    tokenizer, model = load_model()

# Define emotions
emotions = ['anger', 'disgust', 'fear', 'joy', 'neutral', 'sadness', 'surprise']

# Text input
text_input = st.text_area("Enter your text here:", height=150)

# Add example button
if st.button("Try an example"):
    text_input = "I just got the best news ever! I'm so excited and happy I can hardly contain myself! 🎉"
    st.text_area("Enter your text here:", value=text_input, height=150)

if st.button("Analyze Emotions"):
    if text_input.strip() == "":
        st.warning("Please enter some text to analyze.")
    else:
        with st.spinner("Analyzing emotions..."):
            # Get predictions
            probs = analyze_text(text_input, tokenizer, model)
            emotions_dict = dict(zip(emotions, probs))
            
            # Display results
            st.subheader("Analysis Results")
            
            # Create columns for layout
            col1, col2 = st.columns([2, 1])
            
            with col1:
                # Display plot
                fig = create_emotion_plot(emotions_dict)
                st.plotly_chart(fig, use_container_width=True)
            
            with col2:
                # Display scores
                st.subheader("Emotion Scores:")
                for emotion, score in emotions_dict.items():
                    st.write(f"{emotion.capitalize()}: {score:.2%}")

# Add footer
st.markdown("---")
st.markdown("""
Created with ❤️ using Hugging Face Transformers and Streamlit.
Model: j-hartmann/emotion-english-distilroberta-base
""")