#!/bin/bash
set -e

# Remove old CUDA GPG key if present
sudo apt-key del 7fa2af80 2>/dev/null || true

# Install prerequisites
sudo apt-get update
sudo apt-get install -y gcc g++ make dkms

# Add NVIDIA CUDA repository (Ubuntu 22.04/24.04)
DISTRO=$(lsb_release -rs | tr -d '.')  # e.g., 2204 or 2404
ARCH=$(uname -m)  # x86_64 or aarch64

# Download and install the CUDA keyring package
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${DISTRO}/${ARCH}/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
rm cuda-keyring_1.1-1_all.deb

# Update and install the latest CUDA toolkit
sudo apt-get update
sudo apt-get install -y cuda-toolkit   # toolkit only (no driver)
# OR install everything including the driver:
# sudo apt-get install -y cuda

# Add CUDA to PATH and LD_LIBRARY_PATH
CUDA_PATH="/usr/local/cuda"
echo "" >> ~/.bashrc
echo "# CUDA" >> ~/.bashrc
echo "export PATH=${CUDA_PATH}/bin:\$PATH" >> ~/.bashrc
echo "export LD_LIBRARY_PATH=${CUDA_PATH}/lib64:\$LD_LIBRARY_PATH" >> ~/.bashrc
source ~/.bashrc

echo "CUDA installed successfully!"
nvcc --version