import os

# Set the environment variable to point to the CUDA driver in WSL
# This is a common fix for Numba in WSL environments.
if os.path.exists("/usr/lib/wsl/lib"):
    print("WSL detected. Setting NUMBA_CUDA_DRIVER environment variable.")
    os.environ['NUMBA_CUDA_DRIVER'] = '/usr/lib/wsl/lib/libcuda.so.1'

from numba import cuda

print("Checking for Numba CUDA GPU...")
try:
    gpus = cuda.gpus
    if not gpus:
        print("No GPUs found by Numba.")
    else:
        print("Numba found the following GPUs:")
        for i, gpu in enumerate(gpus):
            with gpu:
                print(f"  GPU {i}: {cuda.current_context().device.name.decode('UTF-8')}")
    print("SUCCESS: Numba can access the GPU.")
except Exception as e:
    print(f"ERROR: An exception occurred while trying to access GPUs with Numba.")
    print(e)