"""
Quick test to verify the Kolam AI Generator project is working.
"""

import torch
import numpy as np
import sys
from pathlib import Path

# Add project paths
sys.path.append('models')
sys.path.append('utils')

def test_imports():
    """Test that all modules can be imported."""
    print("🔍 Testing imports...")
    
    try:
        from models.cnn_feature_extractor import KolamFeatureExtractor
        print("✅ CNN Feature Extractor imported successfully")
        
        from models.gan_generator import KolamGenerator
        print("✅ GAN Generator imported successfully")
        
        from utils.image_utils import create_synthetic_kolam
        print("✅ Image utilities imported successfully")
        
        from utils.metrics import KolamDesignMetrics
        print("✅ Metrics imported successfully")
        
        return True
    except Exception as e:
        print(f"❌ Import error: {e}")
        return False

def test_models():
    """Test model creation and basic functionality."""
    print("\n🧠 Testing model creation...")
    
    try:
        from models.cnn_feature_extractor import KolamFeatureExtractor
        from models.gan_generator import KolamGenerator
        
        # Test CNN Feature Extractor
        feature_extractor = KolamFeatureExtractor(input_channels=1, feature_dim=128)
        print("✅ CNN Feature Extractor created successfully")
        
        # Test GAN Generator
        generator = KolamGenerator(noise_dim=100, feature_dim=128, output_channels=1, image_size=64)
        print("✅ GAN Generator created successfully")
        
        # Test basic forward pass
        test_input = torch.randn(2, 1, 64, 64)
        features = feature_extractor(test_input)
        print(f"✅ CNN forward pass: {test_input.shape} → {features.shape}")
        
        noise = torch.randn(2, 100)
        generated = generator(noise)
        print(f"✅ Generator forward pass: {noise.shape} → {generated.shape}")
        
        return True
    except Exception as e:
        print(f"❌ Model test error: {e}")
        return False

def test_utilities():
    """Test utility functions."""
    print("\n🛠️ Testing utilities...")
    
    try:
        from utils.image_utils import create_synthetic_kolam
        from utils.metrics import KolamDesignMetrics
        
        # Test synthetic Kolam creation
        kolam = create_synthetic_kolam(size=(64, 64), complexity='medium')
        print(f"✅ Synthetic Kolam created: {kolam.shape}")
        
        # Test metrics
        metrics = KolamDesignMetrics()
        quality = metrics.calculate_overall_quality(kolam)
        print(f"✅ Quality metrics calculated: {quality['overall_quality']:.4f}")
        
        return True
    except Exception as e:
        print(f"❌ Utility test error: {e}")
        return False

def create_sample_data():
    """Create sample data for demonstration."""
    print("\n📁 Creating sample data...")
    
    try:
        from utils.image_utils import create_synthetic_kolam
        
        # Create directories
        data_dir = Path('data/processed')
        data_dir.mkdir(parents=True, exist_ok=True)
        
        # Create sample images
        for i in range(5):
            kolam = create_synthetic_kolam(size=(64, 64), complexity='medium')
            np.save(data_dir / f'sample_{i:03d}.npy', kolam)
        
        print(f"✅ Created 5 sample images in {data_dir}")
        return True
    except Exception as e:
        print(f"❌ Data creation error: {e}")
        return False

def main():
    """Main test function."""
    print("🚀 Kolam AI Generator - Project Test")
    print("=" * 50)
    
    tests = [
        ("Import Test", test_imports),
        ("Model Test", test_models),
        ("Utility Test", test_utilities),
        ("Data Creation", create_sample_data)
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        print(f"\n📋 Running {test_name}...")
        if test_func():
            passed += 1
            print(f"✅ {test_name} PASSED")
        else:
            print(f"❌ {test_name} FAILED")
    
    print(f"\n📊 Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("\n🎉 All tests passed! The Kolam AI Generator project is ready to use!")
        print("\n🚀 Next steps:")
        print("1. Add real Kolam images to data/raw/")
        print("2. Run: python scripts/preprocess.py")
        print("3. Start training with Jupyter notebooks in notebooks/")
        print("4. Or run training scripts directly")
    else:
        print(f"\n⚠️ {total - passed} tests failed. Please check the errors above.")

if __name__ == "__main__":
    main()