import gradio as gr
from transformers import pipeline
import numpy as np

# Initialize the sentiment classifier
classifier = pipeline(
    "text-classification",
    model="SamLowe/roberta-base-go_emotions",
    top_k=None  # Return all emotions
)

def process_emotions(emotions):
    """Convert raw emotions to our application's format"""
    # Convert to dictionary format
    emotion_dict = {item['label']: float(item['score']) for item in emotions}
    
    # Find primary emotion
    primary_emotion = max(emotion_dict.items(), key=lambda x: x[1])
    
    # Calculate intensity
    max_score = primary_emotion[1]
    intensity = 'High' if max_score > 0.66 else 'Medium' if max_score > 0.33 else 'Low'
    
    # Check if needs attention
    needs_attention = (
        primary_emotion[0] in ['anger', 'anxiety', 'depression'] 
        and max_score > 0.5
    )
    
    return {
        'emotions': emotion_dict,
        'primaryEmotion': primary_emotion[0],
        'emotionalState': {
            'state': primary_emotion[0],
            'intensity': intensity,
            'needsAttention': needs_attention,
            'description': f"Detected {primary_emotion[0]} with {round(max_score * 100)}% confidence"
        },
        'success': True,
        'needsAttention': needs_attention
    }

def analyze_text(text):
    """Analyze text and return processed emotions"""
    if not text or not text.strip():
        return {
            'error': 'Please provide some text to analyze',
            'success': False
        }
    
    try:
        # Get raw emotions from model
        emotions = classifier(text)
        
        # Process emotions into our format
        result = process_emotions(emotions[0])
        
        return result
    except Exception as e:
        return {
            'error': str(e),
            'success': False
        }

# Create Gradio interface
demo = gr.Interface(
    fn=analyze_text,
    inputs=gr.Textbox(label="Enter text to analyze", lines=3),
    outputs=gr.JSON(label="Sentiment Analysis Results"),
    title="Mental Health Sentiment Analysis",
    description="Analyzes text for emotions related to mental health using the RoBERTa model.",
    examples=[
        ["I'm feeling really anxious about my upcoming presentation"],
        ["Today was a great day, I accomplished all my goals!"],
        ["I've been feeling down and unmotivated lately"],
    ],
    allow_flagging="never"
)

# Launch the app
if __name__ == "__main__":
    demo.launch()