#!/usr/bin/env python3
"""
Demo script for Vietnamese Sentiment Analysis
Shows how to use the fine-tuned model for real-time sentiment analysis
"""

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import time

class SentimentDemo:
    def __init__(self, model_path="./vietnamese_sentiment_finetuned"):
        self.model_path = model_path
        self.tokenizer = None
        self.model = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.sentiment_labels = ["Negative", "Neutral", "Positive"]
        
    def load_model(self):
        """Load the fine-tuned model"""
        print(f"🤖 Loading model from: {self.model_path}")
        print(f"📱 Device: {self.device}")
        
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
            self.model = AutoModelForSequenceClassification.from_pretrained(self.model_path)
            self.model.to(self.device)
            self.model.eval()
            print("✅ Model loaded successfully!")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            print("Please run the training first: python run_training.py")
            return False
        
        return True
    
    def predict_sentiment(self, text):
        """Predict sentiment for given text"""
        start_time = time.time()
        
        # Tokenize
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=512
        )
        
        # Move to device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Predict
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=-1)
            predicted_class = torch.argmax(probabilities, dim=-1).item()
            confidence = torch.max(probabilities).item()
        
        inference_time = time.time() - start_time
        
        return {
            "text": text,
            "sentiment": self.sentiment_labels[predicted_class],
            "sentiment_id": predicted_class,
            "confidence": confidence,
            "probabilities": probabilities.cpu().numpy()[0].tolist(),
            "inference_time": inference_time
        }
    
    def demo_mode(self):
        """Run interactive demo"""
        print("\n" + "="*60)
        print("🎭 VIETNAMESE SENTIMENT ANALYSIS DEMO")
        print("="*60)
        print("\n💡 Type Vietnamese text to analyze sentiment")
        print("📝 Type 'quit' to exit, 'help' for examples")
        print("-"*60)
        
        examples = [
            "Giảng viên dạy rất hay và tâm huyết.",
            "Môn học này quá khó và nhàm chán.",
            "Lớp học ổn định, không có gì đặc biệt.",
            "Tôi rất thích cách giảng dạy của thầy cô.",
            "Chương trình học cần cải thiện nhiều."
        ]
        
        while True:
            text = input("\n🔤 Enter text: ").strip()
            
            if text.lower() in ['quit', 'exit', 'q']:
                print("\n👋 Goodbye!")
                break
            
            if text.lower() == 'help':
                print("\n📚 Example texts you can try:")
                for i, example in enumerate(examples, 1):
                    print(f"   {i}. {example}")
                continue
            
            if not text:
                continue
            
            # Make prediction
            result = self.predict_sentiment(text)
            
            # Display result
            sentiment_emoji = {"Negative": "😞", "Neutral": "😐", "Positive": "😊"}
            emoji = sentiment_emoji[result["sentiment"]]
            
            print(f"\n{emoji} Result:")
            print(f"   📝 Text: {result['text']}")
            print(f"   🎯 Sentiment: {result['sentiment']} (Class {result['sentiment_id']})")
            print(f"   📊 Confidence: {result['confidence']:.3f}")
            print(f"   ⏱️  Time: {result['inference_time']:.3f}s")
            
            # Show probability distribution
            print(f"   📈 Probabilities:")
            for i, (label, prob) in enumerate(zip(self.sentiment_labels, result['probabilities'])):
                bar_length = int(prob * 20)
                bar = "█" * bar_length + "░" * (20 - bar_length)
                print(f"      {label}: {bar} {prob:.3f}")
    
    def batch_demo(self):
        """Demo with batch processing"""
        print("\n" + "="*60)
        print("📊 BATCH PROCESSING DEMO")
        print("="*60)
        
        test_texts = [
            "Giảng viên dạy rất hay và tâm huyết.",
            "Môn học này quá khó và nhàm chán.",
            "Lớp học ổn định, không có gì đặc biệt.",
            "Tôi rất thích cách giảng dạy của thầy cô.",
            "Chương trình học cần cải thiện nhiều.",
            "Thời gian biểu hợp lý, dễ theo kịp.",
            "Bài tập quá nhiều và khó.",
            "Môi trường học tập tốt, bạn bè thân thiện."
        ]
        
        print(f"\n📝 Processing {len(test_texts)} texts...")
        
        start_time = time.time()
        results = []
        
        for text in test_texts:
            result = self.predict_sentiment(text)
            results.append(result)
        
        total_time = time.time() - start_time
        
        print(f"\n⏱️  Total time: {total_time:.3f}s")
        print(f"📊 Average time per text: {total_time/len(test_texts):.3f}s")
        
        print(f"\n📋 Results:")
        print("-"*60)
        
        sentiment_counts = {"Positive": 0, "Neutral": 0, "Negative": 0}
        
        for i, result in enumerate(results, 1):
            sentiment_emoji = {"Negative": "😞", "Neutral": "😐", "Positive": "😊"}
            emoji = sentiment_emoji[result["sentiment"]]
            
            print(f"{i:2d}. {emoji} {result['sentiment']:8s} ({result['confidence']:.2f}) - {result['text'][:40]}...")
            sentiment_counts[result["sentiment"]] += 1
        
        print(f"\n📈 Summary:")
        for sentiment, count in sentiment_counts.items():
            emoji = {"Positive": "😊", "Neutral": "😐", "Negative": "😞"}[sentiment]
            percentage = (count / len(results)) * 100
            print(f"   {emoji} {sentiment}: {count} ({percentage:.1f}%)")

def main():
    """Main demo function"""
    print("🎯 Vietnamese Sentiment Analysis Demo")
    print("=====================================")
    
    # Initialize demo
    demo = SentimentDemo()
    
    # Load model
    if not demo.load_model():
        return
    
    # Choose demo mode
    print("\n🎮 Choose demo mode:")
    print("   1. Interactive (type your own text)")
    print("   2. Batch processing (predefined examples)")
    
    while True:
        choice = input("\nEnter choice (1 or 2): ").strip()
        
        if choice == "1":
            demo.demo_mode()
            break
        elif choice == "2":
            demo.batch_demo()
            break
        else:
            print("❌ Invalid choice. Please enter 1 or 2.")

if __name__ == "__main__":
    main()