"""Quick script to check CUDA availability and PyTorch installation."""

import sys

print("=" * 80)
print("CUDA/PyTorch GPU Check")
print("=" * 80)

try:
    import torch
    print(f"✓ PyTorch version: {torch.__version__}")
    
    cuda_available = torch.cuda.is_available()
    print(f"✓ CUDA available: {cuda_available}")
    
    if cuda_available:
        print(f"✓ CUDA version: {torch.version.cuda}")
        print(f"✓ cuDNN version: {torch.backends.cudnn.version()}")
        print(f"✓ GPU count: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
            props = torch.cuda.get_device_properties(i)
            print(f"    Memory: {props.total_memory / (1024**3):.1f} GB")
    else:
        print("\n⚠️  CUDA is NOT available!")
        print("\nPossible causes:")
        print("1. PyTorch was installed without CUDA support")
        print("2. NVIDIA drivers not installed")
        print("3. CUDA toolkit not installed")
        print("\nTo fix:")
        print("1. Check NVIDIA drivers: nvidia-smi")
        print("2. Reinstall PyTorch with CUDA:")
        print("   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121")
        print("   (adjust cu121 to match your CUDA version)")
        
except ImportError:
    print("✗ PyTorch not installed")
    sys.exit(1)

print("=" * 80)