import gradio as gr
import plotly.express as px
import pandas as pd
from transformers import pipeline

# Load model
classifier = pipeline(task="text-classification", model="AR04/Senti", top_k=None)

id2label = {
    0: "admiration", 1: "amusement", 2: "anger", 3: "annoyance", 4: "approval",
    5: "caring", 6: "confusion", 7: "curiosity", 8: "desire", 9: "disappointment",
    10: "disapproval", 11: "disgust", 12: "embarrassment", 13: "excitement",
    14: "fear", 15: "gratitude", 16: "grief", 17: "joy", 18: "love", 19: "nervousness",
    20: "optimism", 21: "pride", 22: "realization", 23: "relief", 24: "remorse",
    25: "sadness", 26: "surprise", 27: "neutral"
}

def classify_and_visualize(text):
    outputs = classifier([text])[0]  # list of dicts
    # Convert to DataFrame
    df = pd.DataFrame(outputs)

    # Radar chart (polar plot)
    fig = px.line_polar(
        df,
        r="score",
        theta="label",
        line_close=True,
        title="Emotion Radar",
    )
    fig.update_traces(fill="toself", line=dict(color="blue"))

    return df.to_dict("records"), fig

with gr.Blocks() as demo:
    gr.Markdown("## 🎭 Emotion Radar Chat App")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(placeholder="Type a message...", lines=2)
            submit_btn = gr.Button("Analyze")
        with gr.Column():
            result_json = gr.JSON()
            result_plot = gr.Plot()

    submit_btn.click(fn=classify_and_visualize, inputs=text_input, outputs=[result_json, result_plot])

if __name__ == "__main__":
    demo.launch()