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

# 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"
}

def split_text(text, max_length=512):
    """
    Splits the input text into chunks where each chunk has tokens <= max_length.
    Splitting is done at sentence boundaries to maintain context.
    """
    # Simple sentence splitting based on periods. For more accurate splitting, consider using nltk or spacy.
    sentences = text.split('. ')
    chunks = []
    current_chunk = ""
    
    for sentence in sentences:
        # Add the period back if it was removed during splitting
        if not sentence.endswith('.'):
            sentence += '.'
        
        # Check if adding the sentence exceeds the max_length
        encoded = tokenizer.encode(current_chunk + " " + sentence, truncation=False)
        if len(encoded) > max_length:
            if current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = sentence
            else:
                # Single sentence longer than max_length, force split
                for i in range(0, len(encoded), max_length):
                    chunk_tokens = encoded[i:i + max_length]
                    chunk_text = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
                    chunks.append(chunk_text.strip())
                current_chunk = ""
        else:
            current_chunk += " " + sentence
    
    if current_chunk:
        chunks.append(current_chunk.strip())
    
    return chunks

@spaces.GPU
def analyze_sentiment(text, show_probabilities=False):
    """
    Analyzes the sentiment of the input text. If the text exceeds the token limit,
    it splits the text into chunks and aggregates the results.
    """
    try:
        chunks = split_text(text)
        all_probabilities = []
        all_predictions = []
        detailed_results = ""

        for idx, chunk in enumerate(chunks, 1):
            inputs = tokenizer(chunk, 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]
            
            all_probabilities.append(probabilities)
            all_predictions.append(predicted_class)
            
            if show_probabilities:
                detailed_results += f"**Chunk {idx}: {predicted_sentiment} ({confidence:.2%})**\n"
                for cls, prob in zip(SENTIMENT_MAP.values(), probabilities):
                    detailed_results += f"{cls}: {prob:.2%}\n"
                detailed_results += "\n"
            else:
                detailed_results += f"**Chunk {idx}: {predicted_sentiment} ({confidence:.2%})**\n"
        
        # Aggregate results by averaging probabilities
        avg_probabilities = sum(all_probabilities) / len(all_probabilities)
        final_class = avg_probabilities.argmax()
        final_sentiment = SENTIMENT_MAP[final_class]
        final_confidence = avg_probabilities[final_class]
        
        result = f"**Overall Sentiment: {final_sentiment}**\nConfidence: {final_confidence:.2%}\n\n"
        
        if show_probabilities:
            result += "### Detailed Analysis:\n" + detailed_results
        else:
            result += "### Detailed Analysis:\n" + detailed_results
        
        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():
            input_text = gr.Textbox(
                lines=10,
                placeholder="Enter text for sentiment analysis...",
                label="Input Text"
            )
            show_probs = gr.Checkbox(
                label="Show probabilities for each class",
                value=False
            )
            analyze_button = gr.Button("Analyze Sentiment")
        with gr.Column():
            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="Predefined Examples"
        )
    
    analyze_button.click(
        fn=analyze_sentiment,
        inputs=[input_text, show_probs],
        outputs=output
    )
    
    gr.Markdown(
        """
        ---
        **Developed with ❤️ using Gradio and Transformers by Hugging Face**
        """
    )

# Launch the interface
demo.launch()