import gradio as gr
import pandas as pd
from transformers import pipeline
import warnings
import os
warnings.filterwarnings("ignore")

# Initialize the models
print("Loading ABSA models for Hugging Face Spaces...")
token_classifier = pipeline(
    model="sdf299/abte-restaurants-distilbert-base-uncased",
    aggregation_strategy="simple"
)

classifier = pipeline(
    model="sdf299/absa-restaurants-distilbert-base-uncased"
)
print("Models loaded successfully!")

def get_sentiment_color(sentiment_label):
    """Return color based on sentiment label."""
    sentiment_lower = sentiment_label.lower()
    if 'positive' in sentiment_lower:
        return "#28a745", "🟢"  # Green
    elif 'negative' in sentiment_lower:
        return "#dc3545", "🔴"  # Red
    else:
        return "#6c757d", "⚪"  # Gray for neutral

def analyze_sentiment(sentence):
    """
    Perform aspect-based sentiment analysis on the input sentence.
    
    Args:
        sentence (str): Input sentence to analyze
        
    Returns:
        tuple: (formatted_results, aspects_summary, detailed_dataframe)
    """
    if not sentence.strip():
        return "Please enter a sentence to analyze.", "", pd.DataFrame()
    
    try:
        # Extract aspects using token classifier
        results = token_classifier(sentence)
        
        if not results:
            return "No aspects found in the sentence.", "", pd.DataFrame()
        
        # Get unique aspects
        aspects = list(set([result['word'] for result in results]))
        
        # Analyze sentiment for each aspect
        detailed_results = []
        formatted_output = f"**📝 Input Sentence:** {sentence}\n\n"
        formatted_output += "## 🎯 Analysis Results:\n\n"
        
        # Count sentiments for summary
        sentiment_counts = {'positive': 0, 'negative': 0, 'neutral': 0}
        
        for aspect in aspects:
            # Classify sentiment for this aspect
            sentiment_result = classifier(f'{sentence} [SEP] {aspect}')
            
            # Extract sentiment label and confidence
            sentiment_label = sentiment_result[0]['label']
            confidence = sentiment_result[0]['score']
            
            # Get color and emoji for this sentiment
            color, emoji = get_sentiment_color(sentiment_label)
            
            # Count sentiments
            if 'positive' in sentiment_label.lower():
                sentiment_counts['positive'] += 1
            elif 'negative' in sentiment_label.lower():
                sentiment_counts['negative'] += 1
            else:
                sentiment_counts['neutral'] += 1
            
            # Format the result with colors
            formatted_output += f'<div style="margin: 15px 0; padding: 15px; border-left: 4px solid {color}; background-color: {color}15; border-radius: 5px;">'
            formatted_output += f'<strong style="color: {color};">{emoji} Aspect: {aspect}</strong><br>'
            formatted_output += f'<span style="color: {color}; font-weight: bold;">Sentiment: {sentiment_label}</span> '
            formatted_output += f'<span style="color: #666; font-size: 0.9em;">(Confidence: {confidence:.3f})</span>'
            formatted_output += '</div>\n\n'
            
            # Store for dataframe with colored styling
            detailed_results.append({
                'Aspect': aspect,
                'Sentiment': sentiment_label,
                'Confidence': f"{confidence:.3f}",
                'Color': color,
                'Emoji': emoji
            })
        
        # Create colorful summary
        aspects_summary = "## 📊 Summary:\n\n"
        aspects_summary += f"**🔍 Total Aspects Found:** {len(aspects)}\n\n"
        
        # Add sentiment breakdown
        if sentiment_counts['positive'] > 0:
            aspects_summary += f"🟢 **Positive:** {sentiment_counts['positive']} aspects\n\n"
        if sentiment_counts['negative'] > 0:
            aspects_summary += f"🔴 **Negative:** {sentiment_counts['negative']} aspects\n\n"
        if sentiment_counts['neutral'] > 0:
            aspects_summary += f"⚪ **Neutral:** {sentiment_counts['neutral']} aspects\n\n"
        
        aspects_summary += f"**📝 Identified Aspects:** {', '.join(aspects)}"
        
        # Create dataframe for tabular view (simplified for table)
        df_data = []
        for result in detailed_results:
            df_data.append({
                'Aspect': result['Aspect'],
                'Sentiment': f"{result['Emoji']} {result['Sentiment']}",
                'Confidence': result['Confidence']
            })
        df = pd.DataFrame(df_data)
        
        return formatted_output, aspects_summary, df
        
    except Exception as e:
        error_msg = f"❌ **Error during analysis:** {str(e)}\n\nPlease try again with a different sentence."
        return error_msg, "", pd.DataFrame()

# Create the Gradio interface
with gr.Blocks(
    title="🍽️ Restaurant Review Analyzer - ABSA",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Arial', sans-serif;
        max-width: 1200px;
    }
    .main-header {
        text-align: center;
        margin-bottom: 30px;
    }
    .sentiment-positive {
        color: #28a745 !important;
        background-color: #d4edda;
        border-color: #c3e6cb;
    }
    .sentiment-negative {
        color: #dc3545 !important;
        background-color: #f8d7da;
        border-color: #f5c6cb;
    }
    .sentiment-neutral {
        color: #6c757d !important;
        background-color: #f8f9fa;
        border-color: #dee2e6;
    }
    """
) as demo:
    
    gr.HTML("""
    <div class="main-header">
        <h1>🍽️ Restaurant Review Analyzer</h1>
        <h3>🎨 Colorful Aspect-Based Sentiment Analysis</h3>
        <p>Analyze restaurant reviews to identify specific aspects and their sentiments with beautiful color coding!</p>
        <div style="margin: 15px 0; padding: 10px; background-color: #f8f9fa; border-radius: 8px; border: 1px solid #dee2e6;">
            <p style="margin: 5px 0;"><strong>🎨 Color Guide:</strong></p>
            <span style="color: #28a745; font-weight: bold;">🟢 Positive</span> | 
            <span style="color: #dc3545; font-weight: bold;">🔴 Negative</span> | 
            <span style="color: #6c757d; font-weight: bold;">⚪ Neutral</span>
        </div>
        <p><em>Powered by DistilBERT models fine-tuned on restaurant reviews</em></p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Input section
            sentence_input = gr.Textbox(
                label="🍽️ Enter Restaurant Review",
                placeholder="e.g., The services here is wonderful, but I hate the food. However, I still love the atmosphere here.",
                lines=3,
                max_lines=5
            )
            
            analyze_btn = gr.Button("🔍 Analyze Sentiment", variant="primary", size="lg")
            
            # Example sentences
            gr.Examples(
                examples=[
                    ["The services here is wonderful, but I hate the food. However, I still love the atmosphere here."],
                    ["The food was amazing and the staff was very friendly, but the restaurant was too noisy."],
                    ["Great location and delicious pizza, but the service was slow and the prices are too high."],
                    ["The ambiance is perfect for a romantic dinner, excellent wine selection, but the dessert was disappointing."],
                    ["Fast service and good value for money, but the food quality could be better."],
                    ["Excellent sushi and attentive waiters, though the wait time was quite long."],
                    ["Beautiful decor and reasonable prices, but the pasta was overcooked."],
                    ["Outstanding customer service and fresh ingredients, highly recommend this place!"],
                    ["Terrible experience - rude staff, cold food, and dirty tables. Never coming back."]
                ],
                inputs=sentence_input,
                label="💡 Try these examples:"
            )
        
        with gr.Column(scale=3):
            # Output section
            with gr.Tab("🎨 Colorful Results"):
                results_output = gr.HTML(label="Visual Analysis Results")
            
            with gr.Tab("📊 Summary Dashboard"):
                aspects_output = gr.Markdown(label="Quick Summary")
            
            with gr.Tab("📈 Data Table"):
                table_output = gr.Dataframe(
                    label="Results Table",
                    headers=["Aspect", "Sentiment", "Confidence"]
                )
    
    # Event handlers
    analyze_btn.click(
        fn=analyze_sentiment,
        inputs=[sentence_input],
        outputs=[results_output, aspects_output, table_output]
    )
    
    sentence_input.submit(
        fn=analyze_sentiment,
        inputs=[sentence_input],
        outputs=[results_output, aspects_output, table_output]
    )
    
    # Footer with model information
    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; padding: 20px; border-top: 1px solid #eee;">
        <p><strong>🤖 Models Used:</strong></p>
        <p>🔤 Aspect Extraction: <a href="https://huggingface.co/sdf299/abte-restaurants-distilbert-base-uncased" target="_blank">sdf299/abte-restaurants-distilbert-base-uncased</a></p>
        <p>😊 Sentiment Classification: <a href="https://huggingface.co/sdf299/absa-restaurants-distilbert-base-uncased" target="_blank">sdf299/absa-restaurants-distilbert-base-uncased</a></p>
        <p style="margin-top: 15px; font-size: 0.9em; color: #666;">
            ✨ This app demonstrates colorful aspect-based sentiment analysis for restaurant reviews using fine-tuned DistilBERT models.
        </p>
    </div>
    """)

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