"""
Test script for the QuickDraw Classifier
Run this to test the classification functionality locally
"""

import os
import sys
import torch
from PIL import Image, ImageDraw
import numpy as np

# Add current directory to path to import app modules
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    from app import classify_drawing, load_categories, CATEGORIES
    print("✅ Successfully imported app modules")
except ImportError as e:
    print(f"❌ Error importing app modules: {e}")
    print("Make sure you've installed the requirements: pip install -r requirements.txt")
    sys.exit(1)

def create_test_drawing(category="cat", size=(400, 400)):
    """
    Create a simple test drawing for testing
    """
    # Create a white canvas
    img = Image.new('RGB', size, 'white')
    draw = ImageDraw.Draw(img)
    
    # Draw simple shapes based on category
    if category == "cat":
        # Draw a simple cat face
        # Head circle
        draw.ellipse([150, 150, 250, 250], outline='black', width=3)
        # Ears
        draw.polygon([(140, 150), (160, 120), (180, 150)], outline='black', width=2)
        draw.polygon([(220, 150), (240, 120), (260, 150)], outline='black', width=2)
        # Eyes
        draw.ellipse([170, 170, 180, 180], fill='black')
        draw.ellipse([220, 170, 230, 180], fill='black')
        # Nose
        draw.polygon([(195, 190), (205, 190), (200, 200)], fill='black')
        # Mouth
        draw.arc([185, 200, 215, 220], 0, 180, fill='black', width=2)
        
    elif category == "house":
        # Draw a simple house
        # Base rectangle
        draw.rectangle([150, 200, 250, 300], outline='black', width=3)
        # Roof triangle
        draw.polygon([(140, 200), (200, 150), (260, 200)], outline='black', width=3)
        # Door
        draw.rectangle([180, 240, 220, 300], outline='black', width=2)
        # Windows
        draw.rectangle([160, 220, 180, 240], outline='black', width=2)
        draw.rectangle([220, 220, 240, 240], outline='black', width=2)
        
    elif category == "car":
        # Draw a simple car
        # Body
        draw.rectangle([120, 200, 280, 240], outline='black', width=3)
        # Roof
        draw.rectangle([150, 180, 250, 200], outline='black', width=3)
        # Wheels
        draw.ellipse([140, 230, 170, 260], outline='black', width=3)
        draw.ellipse([230, 230, 260, 260], outline='black', width=3)
        # Windows
        draw.rectangle([160, 185, 190, 195], outline='black', width=2)
        draw.rectangle([210, 185, 240, 195], outline='black', width=2)
    
    return img

def test_classification():
    """
    Test the classification functionality
    """
    print("\n🧪 Testing QuickDraw Classification")
    print("="*50)
    
    # Test 1: Check if categories are loaded
    print(f"📋 Loaded {len(CATEGORIES)} categories")
    print(f"First 10 categories: {CATEGORIES[:10]}")
    
    # Test 2: Test with simple drawings
    test_categories = ["cat", "house", "car"]
    
    for category in test_categories:
        print(f"\n🎨 Testing with {category} drawing...")
        
        # Create test drawing
        test_img = create_test_drawing(category)
        
        # Classify the drawing
        results = classify_drawing(test_img, top_k=5)
        
        if results:
            print(f"✅ Classification successful!")
            print(f"Top predictions:")
            for i, (pred_category, confidence) in enumerate(results, 1):
                percentage = confidence * 100
                marker = "🎯" if pred_category == category else "  "
                print(f"{marker} {i}. {pred_category.title()} - {percentage:.1f}%")
        else:
            print(f"❌ Classification failed for {category}")
    
    return True

def test_model_loading():
    """
    Test if CLIP model loads correctly
    """
    print("\n🤖 Testing Model Loading")
    print("="*30)
    
    try:
        import clip
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"📱 Using device: {device}")
        
        model, preprocess = clip.load("ViT-B/32", device=device)
        print("✅ CLIP model loaded successfully")
        
        # Test text encoding
        text_inputs = clip.tokenize(["a drawing of a cat", "a drawing of a dog"]).to(device)
        with torch.no_grad():
            text_features = model.encode_text(text_inputs)
        print(f"✅ Text encoding working - shape: {text_features.shape}")
        
        return True
        
    except Exception as e:
        print(f"❌ Model loading failed: {e}")
        return False

def main():
    """
    Run all tests
    """
    print("🚀 QuickDraw Classifier Test Suite")
    print("="*40)
    
    # Test 1: Model loading
    if not test_model_loading():
        print("\n❌ Model loading test failed. Please check your installation.")
        return False
    
    # Test 2: Classification
    if not test_classification():
        print("\n❌ Classification test failed.")
        return False
    
    print("\n" + "="*50)
    print("🎉 All tests passed! Your QuickDraw Classifier is ready!")
    print("\n📖 Next steps:")
    print("1. Run 'python app.py' to start the local app")
    print("2. Deploy to Hugging Face Spaces following the README")
    print("3. Test the deployed app with real drawings")
    
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)