#!/bin/bash

# Define GPU drivers array
gpu_drivers=(
    "A100" "A4000" "A40" "A5000" "A6000" "GT710" "GT730" "GTX1650"
    "GTX1660" "K40" "K620" "K80" "P1000" "P600" "P620" "RTX2060"
    "RTX3060TI" "RTX4060" "RTX4090" "RTXA4000" "RTXA5000" "RTXA6000"
    "T1000" "V100" "P100" "H100"
)

#gpu_drivers=(
#    "a100"
#    "a4000"
#    "a40"
#    "a5000"
#    "a6000"
#    "gt710"
#    "gt730"
#    "gtx1650"
#    "gtx1660"
#    "k40"
#    "k620"
#    "k80"
#    "p1000"
#    "p600"
#    "p620"
#    "rtx2060"
#    "rtx3060ti"
#    "rtx4060"
#    "rtx4090"
#    "rtxa4000"
#    "rtxa5000"
#    "rtxa6000"
#    "t1000"
#    "v100"
#)

# Loop through each GPU driver name
for driver_name in "${gpu_drivers[@]}"; do
    # Create a new script file for each driver
    script_filename_string="ubuntu20-gpu-${driver_name}.sh"
    script_filename=$(echo $script_filename_string | tr '[:upper:]' '[:lower:]')

    # Write the script content to the file
    cat <<EOF > "${script_filename}"
#!/bin/bash

apt update -y &>/dev/null
apt install linux-headers-\$(uname -r) -y &>/dev/null

# Download the GPU driver
wget "https://download.cloudclusters.io/${driver_name}.run" -O ${driver_name}.run &>/dev/null
chmod 755 ${driver_name}.run
./${driver_name}.run -s &>/dev/null

# Function to check if the driver is installed
check_driver() {
    # Use the nvidia-smi command to get the list of GPU devices and filter out lines containing UUID using grep
    gpu_info=\$(nvidia-smi -L | grep -i uuid)

    # Check if GPU information was successfully obtained
    if [ -n "\$gpu_info" ]; then
        echo "Nvidia GPU driver is installed successfully"
        return 0  # Return success status
    else
        echo "Nvidia GPU driver is not installed or installation failed"
        return 1  # Return failure status
    fi
}

check_driver
EOF

done
