# app.py for Hugging Face Space (Gradio)
import gradio as gr
from transformers import pipeline
import plotly.graph_objects as go

# Load pre-trained emotion classifier from Hugging Face
# Model trained on datasets similar to dair-ai/emotion (6 emotions: anger, fear, joy, love, sadness, surprise)
classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)

# Define prediction function for Gradio
def predict_emotion(text):
    # Get model predictions (list of [{label, score}] for each emotion)
    predictions = classifier(text)[0]
    
    # Extract emotion labels and scores
    emotions = [pred['label'] for pred in predictions]
    scores = [pred['score'] for pred in predictions]
    
    # Find the top emotion
    top_emotion = emotions[scores.index(max(scores))]
    
    # Create a bar chart with Plotly
    fig = go.Figure(
        data=[
            go.Bar(x=emotions, y=scores, marker_color=['#FF6384', '#36A2EB', '#FFCE56', '#4BC0C0', '#9966FF', '#FF9F40'])
        ],
        layout=go.Layout(
            title="Emotion Probabilities",
            xaxis_title="Emotions",
            yaxis_title="Probability",
            yaxis_range=[0, 1]
        )
    )
    
    # Return top emotion and chart
    return f"Predicted Emotion: {top_emotion}", fig

# Create Gradio interface
iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(label="Enter a tweet or text", placeholder="e.g., I'm so happy today!"),
    outputs=[
        gr.Text(label="Prediction"),
        gr.Plot(label="Emotion Probabilities")
    ],
    title="Emotion Analyzer",
    description="Enter a tweet or short text to predict its emotion (anger, fear, joy, love, sadness, surprise).",
    examples=[
        ["I'm so excited for the weekend!"],  # Should predict Joy
        ["This news is terrifying."],        # Should predict Fear
        ["I miss you so much."]             # Should predict Sadness
    ]
)

# Launch the interface (handled by Hugging Face Spaces)
if __name__ == "__main__":
    iface.launch()