import gradio as gr
from transformers import pipeline

# Load models
text_emotion_model = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
audio_emotion_model = pipeline("audio-classification", model="superb/wav2vec2-base-superb-er")

# Emoji mapping
emoji_map = {
    "joy": "😊", "happy": "😊", "anger": "😠", "angry": "😠",
    "sadness": "😢", "sad": "😢", "fear": "😨", "calm": "😌",
    "surprise": "😲", "disgust": "🤢", "neutral": "😐", "boredom": "🥱",
    "optimism": "🙂"
}

# Format output
def format_emotion(label, score):
    emoji = emoji_map.get(label.lower(), "❓")
    return f"{emoji} <b>{label.title()}</b> ({int(score * 100)}%)"

# Main prediction function
def detect_emotion(mode, text, audio):
    if mode == "Text" and text:
        result = text_emotion_model(text)[0]
        return f"🗣️ <b>Text Emotion:</b><br>{format_emotion(result['label'], result['score'])}"
    elif mode == "Audio" and audio:
        result = audio_emotion_model(audio)[0]
        return f"🎧 <b>Audio Emotion:</b><br>{format_emotion(result['label'], result['score'])}"
    return "⚠️ <b>Please provide valid input.</b>"

# Clear inputs
def clear_all():
    return "Text", "", None, ""

# Build UI
with gr.Blocks() as demo:
    gr.Markdown("## 🎭 Emotion Detection App (Text or Audio)")

    mode = gr.Radio(["Text", "Audio"], label="Choose Input Mode", value="Text")

    text_input = gr.Textbox(label="💬 Enter your text", visible=True)
    audio_input = gr.Audio(type="filepath", label="🎤 Record or Upload Audio", visible=False)

    output = gr.HTML(label="🎯 Output", value="")

    detect_btn = gr.Button("🔍 Detect Emotion")
    clear_btn = gr.Button("❌ Clear Inputs")

    # Show/hide inputs based on mode
    def toggle_inputs(mode):
        return (
            gr.update(visible=(mode == "Text")),
            gr.update(visible=(mode == "Audio"))
        )
    mode.change(fn=toggle_inputs, inputs=mode, outputs=[text_input, audio_input])

    detect_btn.click(fn=detect_emotion, inputs=[mode, text_input, audio_input], outputs=output)
    clear_btn.click(fn=clear_all, inputs=[], outputs=[mode, text_input, audio_input, output])

demo.launch()