import gradio as gr
import nltk
import plotly.graph_objects as go
from nltk.sentiment.vader import SentimentIntensityAnalyzer

nltk.download("vader_lexicon")
sid = SentimentIntensityAnalyzer()

def sentiment_analysis(text):
    scores = sid.polarity_scores(text)
    del scores["compound"]
    
    labels = list(scores.keys())
    values = list(scores.values())
    colors = ['red' if label == 'neg' else 'green' if label == 'pos' else 'white' for label in labels]
    
    fig = go.Figure(go.Bar(
        x=values,
        y=labels,
        orientation='h',
        marker=dict(color=colors)
    ))
    
    fig.update_layout(
        title="Sentiment Analysis",
        xaxis_title="Score",
        yaxis_title="Sentiment",
        xaxis=dict(range=[0, 1]),
        yaxis=dict(showgrid=True),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    fig.update_xaxes(showgrid=True)
    fig.update_yaxes(showgrid=True)
    
    fig.write_image('sentiment_chart.png')
    return gr.Image('sentiment_chart.png')

demo = gr.Interface(
    fn=sentiment_analysis,
    inputs=gr.Textbox(lines=5, placeholder="Enter a positive or negative sentence here..."),
    outputs="image",
    title="Sentiment Analysis",
    description="Enter a sentence to analyze its sentiment. The model will return the sentiment scores for positive, negative, and neutral tones as a bar chart.",
    examples=[["This is wonderful!"], ["I hate this!"], ["It's okay, not bad."]]
)

demo.launch()