"""
Multilingual Emotion Classifier Usage Example
Author: rmtariq
"""

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

def predict_emotion(text, model, tokenizer):
    """Predict emotion for a given text"""
    
    # Tokenize input
    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
    
    # Get prediction
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(predictions, dim=-1).item()
        confidence = predictions[0][predicted_class].item()
    
    # Get emotion label
    emotion = model.config.id2label[str(predicted_class)]
    
    return emotion, confidence

def main():
    """Main function to demonstrate model usage"""
    
    print("🤖 Multilingual Emotion Classifier Demo")
    print("=" * 50)
    
    # Load model and tokenizer
    print("Loading model...")
    model_name = "rmtariq/multilingual-emotion-classifier"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    
    # Test texts in both languages
    test_texts = [
        # English examples
        "I am absolutely thrilled about this news!",
        "This situation makes me furious!",
        "I'm really worried about the exam tomorrow.",
        "You mean everything to me.",
        "I feel so disappointed and sad.",
        "Wow, I never expected this to happen!",
        
        # Malay examples
        "Saya sangat teruja dengan berita ini!",
        "Keadaan ini membuatkan saya marah!",
        "Saya risau tentang peperiksaan esok.",
        "Awak bermakna segala-galanya bagi saya.",
        "Saya berasa kecewa dan sedih.",
        "Wah, saya tidak sangka ini akan berlaku!"
    ]
    
    # Predict emotions
    for i, text in enumerate(test_texts, 1):
        emotion, confidence = predict_emotion(text, model, tokenizer)
        
        # Determine language
        lang = "🇬🇧 English" if i <= 6 else "🇲🇾 Malay"
        
        print(f"{i:2d}. {lang}")
        print(f"    Text: {text}")
        print(f"    Emotion: {emotion} (confidence: {confidence:.3f})")
        print()
    
    # Interactive mode
    print("\n" + "=" * 50)
    print("Interactive Mode - Enter your own text!")
    print("(Type 'quit' to exit)")
    print("=" * 50)
    
    while True:
        user_text = input("\nEnter text (English or Malay): ").strip()
        
        if user_text.lower() in ['quit', 'exit', 'q']:
            break
        
        if user_text:
            emotion, confidence = predict_emotion(user_text, model, tokenizer)
            print(f"Predicted emotion: {emotion} (confidence: {confidence:.3f})")

if __name__ == "__main__":
    main()