import gradio as gr
from transformers import pipeline
import matplotlib.pyplot as plt
from wordcloud import WordCloud
import json
import os

# Load sentiment analysis model
sentiment_pipeline = pipeline("sentiment-analysis")

# Store last input/output using localStorage
STATE_FILE = "state.json"

def save_state(text, sentiment, confidence):
    with open(STATE_FILE, "w") as f:
        json.dump({"text": text, "sentiment": sentiment, "confidence": confidence}, f)

def load_state():
    if os.path.exists(STATE_FILE):
        with open(STATE_FILE, "r") as f:
            return json.load(f)
    return {"text": "", "sentiment": "", "confidence": ""}

# Function to analyze sentiment in real-time
def analyze_sentiment(text):
    if not text.strip():
        return "Enter text to analyze.", "", None, ""
    
    result = sentiment_pipeline(text)[0]
    sentiment, confidence = result['label'], result['score']
    
    sentiment_map = {
        "POSITIVE": ("🟢 Positive 😊", "green"),
        "NEGATIVE": ("🔴 Negative 😠", "red"),
        "NEUTRAL": ("🟡 Neutral 😐", "orange")
    }
    
    sentiment_label, color = sentiment_map.get(sentiment.upper(), ("⚪ Unknown ❓", "gray"))
    
    # Generate Word Cloud
    wordcloud = WordCloud(width=400, height=200, background_color="white").generate(text)
    fig, ax = plt.subplots()
    ax.imshow(wordcloud, interpolation='bilinear')
    ax.axis("off")
    
    save_state(text, sentiment_label, confidence)
    
    return sentiment_label, f"Confidence: {confidence:.2%}", fig, text

# UI Layout
with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("# 🎭 AI Sentiment Analyzer")
    gr.Markdown("Analyze sentiment in real-time with enhanced visualization.")
    
    with gr.Row():
        dark_mode = gr.Checkbox(label="🌙 Dark Mode")
        
    with gr.Row():
        text_input = gr.Textbox(lines=3, placeholder="Type your text here...", label="Your Input", live=True)
    
    analyze_button = gr.Button("Analyze Sentiment ✨")
    reset_button = gr.Button("🔄 Reset")
    
    with gr.Row():
        sentiment_output = gr.Label(label="Sentiment Result")
        confidence_output = gr.Label(label="Confidence Score")
    
    wordcloud_output = gr.Plot(label="Word Cloud Visualization")
    
    # Load previous session state
    state = load_state()
    text_input.value, sentiment_output.value, confidence_output.value = state['text'], state['sentiment'], state['confidence']
    
    analyze_button.click(analyze_sentiment, inputs=text_input, outputs=[sentiment_output, confidence_output, wordcloud_output, text_input])
    reset_button.click(lambda: ("", "", None, ""), inputs=[], outputs=[sentiment_output, confidence_output, wordcloud_output, text_input])

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