import gradio as gr
from transformers import pipeline
import pandas as pd
import plotly.graph_objects as go

# Load the emotion classification model
model_id = "S-4-G-4-R/distilbert-base-uncased-finetuned-emotion"
classifier = pipeline("text-classification", model=model_id)

# Define emotion labels mapping (LABEL_0 to LABEL_5)
EMOTION_LABELS = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']

# Label mapping from model output to emotion names
LABEL_MAPPING = {
    'LABEL_0': 'sadness',
    'LABEL_1': 'joy',
    'LABEL_2': 'love',
    'LABEL_3': 'anger',
    'LABEL_4': 'fear',
    'LABEL_5': 'surprise'
}

# Emoji mapping for emotions
EMOTION_EMOJIS = {
    'sadness': '😢',
    'joy': '😊',
    'love': '❤️',
    'anger': '😠',
    'fear': '😨',
    'surprise': '😲'
}

def classify_emotion(text):
    """
    Classify the emotion in the given text and return results with visualization
    """
    if not text.strip():
        return None, "Please enter some text to analyze."
    
    # Get predictions
    preds = classifier(text, return_all_scores=True)[0]
    
    # Create DataFrame and map labels to emotion names
    df = pd.DataFrame(preds)
    df['emotion'] = df['label'].map(LABEL_MAPPING)
    df['score'] = df['score'] * 100  # Convert to percentage
    
    # Add emojis to labels
    df['display_label'] = df['emotion'].map(lambda x: f"{EMOTION_EMOJIS.get(x, '')} {x.capitalize()}")
    
    # Sort by score for better visualization
    df = df.sort_values('score', ascending=True)
    
    # Create horizontal bar chart using Plotly
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A', '#98D8C8', '#F7DC6F']
    
    fig = go.Figure(go.Bar(
        x=df['score'],
        y=df['display_label'],
        orientation='h',
        marker=dict(
            color=df['score'],
            colorscale='Viridis',
            showscale=False
        ),
        text=df['score'].round(2),
        texttemplate='%{text}%',
        textposition='outside'
    ))
    
    fig.update_layout(
        title={
            'text': 'Emotion Classification Results',
            'x': 0.5,
            'xanchor': 'center'
        },
        xaxis_title='Confidence (%)',
        yaxis_title='',
        height=450,
        margin=dict(l=20, r=80, t=60, b=40),
        plot_bgcolor='rgba(13, 13, 9, 0.05)',
        paper_bgcolor='rgba(13, 13, 9, 0.05)',
        font=dict(size=12, color='white')
    )
    
    fig.update_xaxes(range=[0, 105], gridcolor='lightgray')
    
    # Format results as text with emojis
    results_text = "### 🎯 Prediction Results\n\n"
    sorted_df = df.sort_values('score', ascending=False)
    
    top_emotion = sorted_df.iloc[0]
    results_text += f"**Top Emotion:** {EMOTION_EMOJIS.get(top_emotion['emotion'], '')} **{top_emotion['emotion'].capitalize()}** ({top_emotion['score']:.2f}%)\n\n"
    results_text += "---\n\n**All Emotions:**\n\n"
    
    for _, row in sorted_df.iterrows():
        emoji = EMOTION_EMOJIS.get(row['emotion'], '')
        bar_length = int(row['score'] / 5)
        bar = '█' * bar_length
        results_text += f"{emoji} **{row['emotion'].capitalize()}**: {row['score']:.2f}% {bar}\n\n"
    
    return fig, results_text

# Example texts
examples = [
    ["I was feeling very alone today walking down on road"],
    ["I am so happy and excited about this new opportunity!"],
    ["This makes me really angry and frustrated!"],
    ["I'm scared about what might happen next..."],
    ["What a beautiful day, I love this!"],
    ["Wow! I can't believe this just happened!"],
    ["I feel so sad and disappointed about the news."]
]

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), title="Emotion Classifier") as demo:
    gr.Markdown(
        """
        # 🎭 Emotion Classification
        Analyze the emotional tone of any text using AI. This model can detect **6 emotions**: 
        Sadness 😢, Joy 😊, Love ❤️, Anger 😠, Fear 😨, and Surprise 😲
        
        **Model:** S-4-G-4-R/distilbert-base-uncased-finetuned-emotion
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            text_input = gr.Textbox(
                label="📝 Enter text to analyze",
                placeholder="Type or paste your text here...",
                lines=5
            )
            classify_btn = gr.Button("🔍 Classify Emotion", variant="primary", size="lg")
            
            gr.Markdown("### 💡 Try these examples:")
            gr.Examples(
                examples=examples,
                inputs=text_input,
                label=None
            )
        
        with gr.Column(scale=1):
            results_text = gr.Markdown(label="Results")
    
    with gr.Row():
        plot_output = gr.Plot(label="📊 Emotion Probabilities")
    
    gr.Markdown(
        """
        ---
        **How it works:** The model analyzes your text and assigns confidence scores to each of the 6 emotions. 
        Higher percentages indicate stronger presence of that emotion in the text.
        """
    )
    
    # Connect the button
    classify_btn.click(
        fn=classify_emotion,
        inputs=text_input,
        outputs=[plot_output, results_text]
    )
    
    # Also trigger on Enter key
    text_input.submit(
        fn=classify_emotion,
        inputs=text_input,
        outputs=[plot_output, results_text]
    )

# Launch the app
demo.launch()