#!/usr/bin/env python3
"""
Test script to verify all dependencies are working correctly
"""

import sys
import traceback

def test_imports():
    """Test importing all required dependencies"""
    print("Testing imports...")
    
    try:
        import torch
        print(f"✓ PyTorch {torch.__version__}")
        print(f"  CUDA available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"  CUDA version: {torch.version.cuda}")
            print(f"  GPU count: {torch.cuda.device_count()}")
    except Exception as e:
        print(f"✗ PyTorch import failed: {e}")
        return False
    
    # Check if torch-sparse is available or disabled
    try:
        import torch_sparse
        print("✓ torch-sparse")
    except ImportError:
        # Check if torch-sparse was disabled
        try:
            with open("NeuralJacobianFields/PoissonSystem.py", 'r') as f:
                content = f.read()
            if "USE_TORCH_SPARSE = False" in content:
                print("✓ torch-sparse (disabled, using built-in PyTorch sparse)")
            else:
                print("✗ torch-sparse (not available)")
                return False
        except:
            print("✗ torch-sparse (not available)")
            return False
    except Exception as e:
        print(f"✗ torch-sparse import failed: {e}")
        return False
    
    # Check if torch-scatter is available (not critical)
    try:
        import torch_scatter
        print("✓ torch-scatter")
    except ImportError:
        print("⚠ torch-scatter (not available, may not be critical)")
    except Exception as e:
        print(f"⚠ torch-scatter import failed: {e} (may not be critical)")
    
    try:
        import nvdiffrast
        print("✓ nvdiffrast")
    except Exception as e:
        print(f"✗ nvdiffrast import failed: {e}")
        return False
    
    try:
        import pytorch3d
        print("✓ PyTorch3D")
    except Exception as e:
        print(f"✗ PyTorch3D import failed: {e}")
        return False
    
    try:
        import clip
        print("✓ CLIP")
    except Exception as e:
        print(f"✗ CLIP import failed: {e}")
        return False
    
    return True

def test_basic_functionality():
    """Test basic functionality of key components"""
    print("\nTesting basic functionality...")
    
    try:
        import torch
        
        # Test basic tensor operations
        x = torch.randn(10, 10)
        y = torch.randn(10, 10)
        z = x + y
        print("✓ Basic tensor operations")
        
        # Test sparse operations (using built-in PyTorch sparse)
        indices = torch.randint(0, 10, (2, 20))
        values = torch.randn(20)
        sparse_tensor = torch.sparse_coo_tensor(indices, values, (10, 10))
        print("✓ Sparse tensor creation")
        
        # Test if torch-sparse is available (optional)
        try:
            import torch_sparse
            print("✓ torch-sparse operations available")
        except ImportError:
            print("✓ Using built-in PyTorch sparse operations")
        
        # Test if torch-scatter is available (optional)
        try:
            import torch_scatter
            print("✓ torch-scatter operations available")
        except ImportError:
            print("✓ torch-scatter not available (not critical)")
        
    except Exception as e:
        print(f"✗ Basic functionality test failed: {e}")
        traceback.print_exc()
        return False
    
    return True

def main():
    """Main test function"""
    print("=== Garment3DGen Dependency Test ===\n")
    
    # Test imports
    if not test_imports():
        print("\n❌ Import tests failed!")
        sys.exit(1)
    
    # Test basic functionality
    if not test_basic_functionality():
        print("\n❌ Functionality tests failed!")
        sys.exit(1)
    
    print("\n✅ All tests passed! Dependencies are working correctly.")
    print("The Garment3DGen application should be ready to run.")

if __name__ == "__main__":
    main()