#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091
## This script aims to setup the OS for a (cloud) VM
## starting from the "Plain Ubuntu 22.04" image
#! Update and upgrade package manager
sudo apt-get update
#! Installing the essentials
sudo apt-get install -y build-essential zlib1g-dev libedit-dev \
	libssl-dev liblzma-dev libffi-dev libbz2-dev \
	libreadline-dev libsqlite3-dev bc clang
#! Check the output of `nvcc -V`
NVCC_OUTPUT=$(nvcc -V)
if [[ $NVCC_OUTPUT == *"release 12.4"* ]]; then
	echo "CUDA 12.4 is detected."
else
	#! Get and install CUDA 12.4.1 and its drivers
	wget wget https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
	sudo sh cuda_12.4.0_550.54.14_linux.run --toolkit --driver --no-man-page --silent
fi
if [[ $PATH == *"cuda-12.4"* ]]; then
	echo "PATH variable is already set."
else
	#! Set the PATH env variables for the new CUDA version
	echo '# Adding CUDA 12.4 to the PATH environmental variables' >>~/.bashrc
	# shellcheck disable=SC2016
	echo 'export PATH=/usr/local/cuda-12.4/bin${PATH:+:${PATH}}' >>~/.bashrc
	export PATH=/usr/local/cuda-12.4/bin${PATH:+:${PATH}}
fi
if [[ $LD_LIBRARY_PATH == *"cuda-12.4"* ]]; then
	echo "LD_LIBRARY_PATH variable is already set."
else
	#! Set the LD_LIBRARY_PATH env variables for the new CUDA version
	# shellcheck disable=SC2016
	echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >>~/.bashrc
	export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
fi
#! Set GPU persistence mode
sudo nvidia-smi -pm 1

#! Install CuDNN latest
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cudnn9-cuda-12 nvidia-fabricmanager-575 libnvidia-nscq-575
sudo systemctl start nvidia-fabricmanager

#! Install `UV`
UV_VER_OUTPUT=$(uv --version)
if [[ $UV_VER_OUTPUT == *"uv "* ]]; then
	echo "UV is already installed."
else
	#! Getting `UV`
	curl -LsSf https://astral.sh/uv/install.sh | sh
fi

#! Upgrade pip
sudo pip install --upgrade pip
#! Monitoring utilities and cmake
sudo apt -y install bpytop
sudo pip install nvitop cmake
