import torch
import os

# Check if PyTorch is installed
print(f"Torch version: {torch.__version__}")

# Check if CUDA (GPU support) is available
cuda_available = torch.cuda.is_available()
print(f"CUDA Available: {cuda_available}")

if not cuda_available:
    print("Try executing in a queue with GPU available (E.g. acc_debug)")

# Perform a basic tensor operation
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])
c = a + b

print(f"Tensor addition result: {c}")

# Check if the summation is correct
expected = torch.tensor([5.0, 7.0, 9.0])
if torch.equal(c, expected):
    print("Summation is correct!")
else:
    print("Summation is incorrect!")

# SLURM specific checks
if "SLURM_GPUS_ON_NODE" in os.environ:
    # Check if the number of GPUs on the node matches what is available to torch
    if int(os.environ["SLURM_GPUS_ON_NODE"]) != torch.cuda.device_count():
        print("Error: Torch should see and use all available GPUs, but it does not.")

    # Check if the CPUs per task set in SLURM matches what PyTorch sees
    if int(os.environ["SLURM_CPUS_PER_TASK"]) != torch.get_num_threads():
        print("Error: in SLURM, --cpus-per-task sets this up, and torch should see it.")

