#!/usr/bin/env python3
"""
Example script to test the Arabic Message Classification Model
"""

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch

def main():
    # Model name - replace with your actual model name on Hugging Face
    model_name = "ahmedmajid92/Arabic_MI_Classifier"
    
    print("Loading Arabic Message Classification Model...")
    print(f"Model: {model_name}")
    
    try:
        # Load tokenizer and model
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSequenceClassification.from_pretrained(model_name)
        
        # Create classification pipeline
        device = 0 if torch.cuda.is_available() else -1
        classifier = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer,
            device=device
        )
        
        print(f"Model loaded successfully!")
        print(f"Using device: {'GPU' if device >= 0 else 'CPU'}")
        print("-" * 50)
        
        # Test examples
        test_examples = [
            "السلام عليكم ورحمة الله وبركاته",  # greeting
            "هلو شلونك اليوم؟",                    # greeting + question
            "متى يبدأ الاجتماع؟",                  # question
            "عندي مشكلة بالانترنت",                # complaint
            "أحب القراءة والكتابة",                # general
            "الكهرباء نفطت",                      # complaint (Iraqi)
            "شنو الأخبار؟",                       # question (Iraqi)
            "تحية طيبة",                         # greeting
            "أعمل مهندساً في شركة تقنية",          # general
            "الطابعة ما تطبع"                     # complaint (Iraqi)
        ]
        
        print("Testing with example messages:")
        print("=" * 60)
        
        for i, text in enumerate(test_examples, 1):
            result = classifier(text)[0]
            label = result['label']
            confidence = result['score']
            
            print(f"{i:2d}. Text: {text}")
            print(f"    → Label: {label}")
            print(f"    → Confidence: {confidence:.4f}")
            print()
        
        print("=" * 60)
        print("Interactive mode - Enter your own text (or 'quit' to exit):")
        
        while True:
            user_input = input("\nEnter Arabic text: ").strip()
            
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("Goodbye!")
                break
                
            if not user_input:
                continue
                
            try:
                result = classifier(user_input)[0]
                label = result['label']
                confidence = result['score']
                
                print(f"→ Label: {label}")
                print(f"→ Confidence: {confidence:.4f}")
                
            except Exception as e:
                print(f"Error processing text: {e}")
                
    except Exception as e:
        print(f"Error loading model: {e}")
        print("Make sure to:")
        print("1. Install required packages: pip install transformers torch")
        print("2. Update the model_name variable with your actual model name")
        print("3. Check your internet connection")

if __name__ == "__main__":
    main()