"""
GPU Utilization Test Script
Run this to verify your GPU is being used properly.
"""

import torch
import time
import numpy as np

def test_gpu():
    """Test GPU availability and performance."""
    
    print("=" * 60)
    print("GPU UTILIZATION TEST")
    print("=" * 60)
    
    # Check CUDA availability
    print(f"\n✓ PyTorch version: {torch.__version__}")
    print(f"✓ CUDA available: {torch.cuda.is_available()}")
    
    if not torch.cuda.is_available():
        print("\n❌ ERROR: CUDA not available!")
        print("   Your GPU won't be utilized.")
        print("   Make sure you have:")
        print("   1. CUDA-compatible GPU")
        print("   2. NVIDIA drivers installed")
        print("   3. PyTorch with CUDA support")
        return
    
    # GPU info
    print(f"\n📊 GPU Information:")
    print(f"   Device count: {torch.cuda.device_count()}")
    print(f"   Current device: {torch.cuda.current_device()}")
    print(f"   Device name: {torch.cuda.get_device_name(0)}")
    
    props = torch.cuda.get_device_properties(0)
    print(f"   Total memory: {props.total_memory / 1024**3:.2f} GB")
    print(f"   CUDA capability: {props.major}.{props.minor}")
    
    # Memory test
    print(f"\n💾 Memory Status:")
    memory_allocated = torch.cuda.memory_allocated(0) / 1024**3
    memory_reserved = torch.cuda.memory_reserved(0) / 1024**3
    print(f"   Allocated: {memory_allocated:.3f} GB")
    print(f"   Reserved: {memory_reserved:.3f} GB")
    
    # Performance test
    print(f"\n🚀 Performance Test:")
    print("   Creating large tensors on GPU...")
    
    # Clear cache first
    torch.cuda.empty_cache()
    
    # Create tensors
    device = torch.device("cuda")
    size = 2000
    
    print(f"   Matrix size: {size}x{size}")
    print("   Running matrix multiplication...")
    
    # Warm-up
    a = torch.randn(size, size, device=device)
    b = torch.randn(size, size, device=device)
    _ = torch.matmul(a, b)
    torch.cuda.synchronize()
    
    # Timed test
    iterations = 10
    start_time = time.time()
    
    for i in range(iterations):
        a = torch.randn(size, size, device=device)
        b = torch.randn(size, size, device=device)
        c = torch.matmul(a, b)
        torch.cuda.synchronize()  # Wait for GPU to finish
        
        if i == 0:
            print(f"   ⚡ GPU is working! Check Task Manager now.")
            print(f"   (GPU utilization should spike to 80-100%)")
    
    end_time = time.time()
    elapsed = end_time - start_time
    
    print(f"\n   Completed {iterations} iterations in {elapsed:.2f}s")
    print(f"   Average: {elapsed/iterations:.3f}s per iteration")
    
    # Memory usage after test
    memory_allocated = torch.cuda.memory_allocated(0) / 1024**3
    memory_reserved = torch.cuda.memory_reserved(0) / 1024**3
    print(f"\n   Memory after test:")
    print(f"   Allocated: {memory_allocated:.3f} GB")
    print(f"   Reserved: {memory_reserved:.3f} GB")
    
    # Load CLIP test
    print(f"\n🎨 Testing CLIP model loading...")
    try:
        import clip
        print("   Loading CLIP ViT-B/32...")
        model, preprocess = clip.load("ViT-B/32", device=device)
        print("   ✓ CLIP loaded successfully on GPU")
        
        # Test inference
        print("   Running test inference...")
        test_image = torch.randn(1, 3, 224, 224, device=device)
        with torch.no_grad():
            features = model.encode_image(test_image)
        torch.cuda.synchronize()
        print("   ✓ CLIP inference working")
        
        memory_allocated = torch.cuda.memory_allocated(0) / 1024**3
        print(f"   Memory with CLIP: {memory_allocated:.3f} GB")
        
    except Exception as e:
        print(f"   ❌ CLIP test failed: {e}")
    
    # Cleanup
    torch.cuda.empty_cache()
    
    print("\n" + "=" * 60)
    print("TEST COMPLETE")
    print("=" * 60)
    print("\nWhat you should see in Task Manager:")
    print("  • GPU 1 (NVIDIA GeForce GTX 1650): 80-100% during test")
    print("  • GPU Memory: ~1-2GB allocated")
    print("  • If utilization stays at 0%, there's a problem!")
    print("\nIf GPU utilization was 0% during this test:")
    print("  1. CUDA drivers may not be installed correctly")
    print("  2. PyTorch may be CPU-only version")
    print("  3. Windows may be using integrated GPU (GPU 0)")
    print("\nRecommendation: Run this script and watch Task Manager!")
    print("=" * 60)


if __name__ == "__main__":
    test_gpu()