import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Initialize device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load model and tokenizer
model_name = "tabularisai/robust-sentiment-analysis"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)

# Define sentiment mapping
SENTIMENT_MAP = {
    0: "Very Negative",
    1: "Negative",
    2: "Neutral",
    3: "Positive",
    4: "Very Positive"
}

@spaces.GPU
def analyze_sentiment(text, show_probabilities=False):
    """
    Analyzes the sentiment of the input text.
    """
    try:
        # Preprocess text - convert to lowercase
        text = text.lower()
        
        # Tokenize and prepare input
        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)

        with torch.no_grad():
            outputs = model(**inputs)

        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
        predicted_class = probabilities.argmax()
        predicted_sentiment = SENTIMENT_MAP[predicted_class]
        confidence = probabilities[predicted_class]

        # Prepare the result with emoji indicators
        sentiment_emojis = {
            "Very Negative": "😡",
            "Negative": "😔",
            "Neutral": "😐",
            "Positive": "😊",
            "Very Positive": "🤩"
        }

        result = f"**{sentiment_emojis[predicted_sentiment]} Overall Sentiment: {predicted_sentiment}**\n"
        result += f"Confidence: {confidence:.2%}\n\n"

        if show_probabilities:
            result += "### Detailed Analysis:\n"
            for cls, prob in zip(SENTIMENT_MAP.values(), probabilities):
                emoji = sentiment_emojis[cls]
                result += f"{emoji} {cls}: {prob:.2%}\n"

        return result
    except Exception as e:
        return f"An error occurred during sentiment analysis: {str(e)}"

# Create Gradio interface using Blocks for better layout control
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎭 Sentiment Analysis Wizard")
    gr.Markdown(
        """
        Discover the emotional tone behind any text with our advanced AI model! This app uses a state-of-the-art language model to analyze the sentiment of your text, classifying it into one of five categories: **Very Negative**, **Negative**, **Neutral**, **Positive**, or **Very Positive**.
        """
    )

    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                lines=10,
                placeholder="Enter text for sentiment analysis...",
                label="✍️ Input Text"
            )
            with gr.Row():
                show_probs = gr.Checkbox(
                    label="🎯 Show probabilities for each class",
                    value=False
                )
                analyze_button = gr.Button("✨ Analyze Sentiment", variant="primary")
        
        with gr.Column(scale=1):
            output = gr.Markdown(label="Result")

    with gr.Accordion("📚 Examples", open=False):
        examples = [
            ["I absolutely loved this movie! The acting was superb and the plot was engaging.", True],
            ["The service at this restaurant was terrible. I'll never go back.", False],
            ["The product works as expected. Nothing special, but it gets the job done.", True],
            ["I'm somewhat disappointed with my purchase. It's not as good as I hoped.", False],
            ["This book changed my life! I couldn't put it down and learned so much.", True]
        ]
        gr.Examples(
            examples=examples,
            inputs=[input_text, show_probs],
            label="Try these examples"
        )

    analyze_button.click(
        fn=analyze_sentiment,
        inputs=[input_text, show_probs],
        outputs=output
    )

    gr.Markdown(
        """
        ---
        **Developed by tabularis.ai with ❤️ using Gradio and Transformers by Hugging Face**
        """
    )

# Launch the interface
demo.launch()