# Author: Siwen Yu (yusiwen@gmail.com)
# License: MIT

import torch
import sys

print(f"Python Version: {sys.version.split()[0]}")
print(f"PyTorch Version: {torch.__version__}")

# 1. Check if CUDA (NVIDIA Driver Connection) is available
cuda_available = torch.cuda.is_available()
print(f"CUDA Available: {cuda_available}")

if cuda_available:
    # 2. Check GPU Name and Device Count
    print(f"GPU Count: {torch.cuda.device_count()}")
    print(f"GPU Name: {torch.cuda.get_device_name(0)}")
    print(f"CUDA Capability: {torch.cuda.get_device_capability(0)}")
    
    # 3. Perform a quick mathematical matrix multiplication on the GPU
    print("\n--- Running GPU Hardware Test ---")
    try:
        # Create two random tensors directly on the GPU memory
        x = torch.randn(1000, 1000, device="cuda")
        y = torch.randn(1000, 1000, device="cuda")
        
        # Multiply them (Matrix Multiplication)
        z = torch.matmul(x, y)
        
        print("SUCCESS: Matrix multiplication completed on GPU successfully!")
        
        # Check current VRAM allocation
        allocated_vram = torch.cuda.memory_allocated(0) / (1024 ** 2)
        print(f"Allocated VRAM during test: {allocated_vram:.2f} MB")
        
    except Exception as e:
        print(f"FAILURE: Failed to compute on GPU. Error: {e}")
else:
    print("\nERROR: PyTorch cannot find your NVIDIA GPU.")
    print("Please check your Windows NVIDIA drivers or WSL2 CUDA toolkit setup.")