#!/bin/bash

#----------- Notes ------------------
#check for NVIDIA device:
#apt-get -y install pciutils
#lspci | grep VGA
#check if driver works with device: nvidia-smi

#based on https://docs.nvidia.com/cuda/cuda-installation-guide-linux/
#and https://www.tensorflow.org/install/gpu
#Versions required: https://www.tensorflow.org/install/source#gpu

#remove cuda:
#sudo apt-get remove cuda-11-2
#sudo apt-get autoremove
#------------------------------------

set -e

if [ "$EUID" -ne 0 ]
  then echo "Please run as root"
  exit
fi

distribution=$(. /etc/os-release;echo $ID$VERSION_ID | sed -e 's/\.//g')
architecture="x86_64"

#echo
#echo "Disable cloud-init..."
#this is optional, not a requirement
#touch /etc/cloud/cloud-init.disabled

echo
echo "Install pre-requirements..."
apt-get update
apt-get install -y linux-headers-$(uname -r) software-properties-common


echo
echo "Setup the CUDA repository public GPG key..."
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/$distribution/$architecture/3bf863cc.pub
wget https://developer.download.nvidia.com/compute/cuda/repos/$distribution/$architecture/cuda-$distribution.pin
mv cuda-$distribution.pin /etc/apt/preferences.d/cuda-repository-pin-600
#add-apt-repository contrib



echo
echo "Setup the CUDA network repository..."
add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/$distribution/$architecture/ /"



echo
echo "Install CUDA..."
#Installs CUDA Toolkit and Driver packages.
#Use the --no-install-recommends option should prevent install any dependencies on X packages
#It still installes alot of dependencies, but fortunately it seams not to setup x-server
apt-get update
apt-get install -y --no-install-recommends cuda-11-2


#See for latest CUDA-Version: https://developer.nvidia.com/cudnn-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04&target_type=deb_network
echo
echo "Install libcudnn..."
apt-get install -y libcudnn8_8.1.1.33-1+cuda11.2 libcudnn8-dev_8.1.1.33-1+cuda11.2

#get an other version:
#wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/libcudnn8_8.1.1.33-1+cuda11.2_amd64.deb
#apt-get install -y ./libcudnn8_8.1.1.33-1+cuda11.2_amd64.deb
#rm ./libcudnn8_8.1.1.33-1+cuda11.2_amd64.deb
#wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/libcudnn8-dev_8.1.1.33-1+cuda11.2_amd64.deb
#apt-get install -y ./libcudnn8-dev_8.1.1.33-1+cuda11.2_amd64.deb
#rm ./libcudnn8-dev_8.1.1.33-1+cuda11.2_amd64.deb


echo
echo "Setup successfully finished"
echo "Type \"sudo systemctl reboot\" to reboot system for loading GPU drivers)"