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

# Initialize the models
print("Loading models...")
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 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**Analysis Results:**\n\n"
        
        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']
            
            # Format the result
            formatted_output += f"🎯 **Aspect:** {aspect}\n"
            formatted_output += f"   **Sentiment:** {sentiment_label} (Confidence: {confidence:.3f})\n\n"
            
            # Store for dataframe
            detailed_results.append({
                'Aspect': aspect,
                'Sentiment': sentiment_label,
                'Confidence': f"{confidence:.3f}"
            })
        
        # Create summary
        aspects_summary = f"**Identified Aspects:** {', '.join(aspects)}"
        
        # Create dataframe for tabular view
        df = pd.DataFrame(detailed_results)
        
        return formatted_output, aspects_summary, df
        
    except Exception as e:
        error_msg = f"Error during analysis: {str(e)}"
        return error_msg, "", pd.DataFrame()

def create_interface():
    """Create and configure the Gradio interface."""
    
    with gr.Blocks(
        title="Aspect-Based Sentiment Analysis",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            font-family: 'Arial', sans-serif;
        }
        .main-header {
            text-align: center;
            margin-bottom: 30px;
        }
        """
    ) as demo:
        
        gr.HTML("""
        <div class="main-header">
            <h1>🍽️ Restaurant Review Analyzer</h1>
            <h3>Aspect-Based Sentiment Analysis</h3>
            <p>Analyze restaurant reviews to identify specific aspects (food, service, atmosphere, etc.) and their associated sentiments.</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."]
                    ],
                    inputs=sentence_input
                )
            
            with gr.Column(scale=3):
                # Output section
                with gr.Tab("📊 Detailed Results"):
                    results_output = gr.Markdown(label="Analysis Results")
                
                with gr.Tab("📋 Quick Summary"):
                    aspects_output = gr.Markdown(label="Aspects 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
        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: <code>sdf299/abte-restaurants-distilbert-base-uncased</code></p>
            <p>😊 Sentiment Classification: <code>sdf299/absa-restaurants-distilbert-base-uncased</code></p>
        </div>
        """)
    
    return demo

if __name__ == "__main__":
    # Create and launch the interface
    demo = create_interface()
    demo.launch(
        share=True,  # Creates a public link
        server_name="0.0.0.0",  # Makes it accessible from other devices on the network
        server_port=7860,
        show_error=True
    )