#!/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
#------------------------------------

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 "\n"
echo "Install drivers..."
apt-get update && \
apt-get install -y linux-headers-$(uname -r) software-properties-common

if [ $? -eq 0 ]; then
  echo "\n"
  echo "Install the CUDA repository public GPG key..."
  apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/$distribution/$architecture/7fa2af80.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
fi

if [ $? -eq 0 ]; then
  echo "\n"
  echo "Setup the CUDA network repository..."
  add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/$distribution/$architecture/ /"
fi

if [ $? -eq 0 ]; then
  echo "\n"
  echo "install CUDA..."
  #Installs all CUDA Toolkit and Driver packages. Handles upgrading to the next version of the cuda package when it's released.
  #Use the --no-install-recommends option for a lean driver install without any dependencies on X packages
  apt-get update && \
  apt-get install -y --no-install-recommends cuda-11-2
  #apt-get install --no-install-recommends libcudnn8=8.1.0.44-1+cuda11.2 libcudnn8-dev=8.1.0.44-1+cuda11.2
fi

if [ $? -eq 0 ]; then
  echo "\n"
  echo "install libcudnn8..."
  #Latest available version from ubuntu1804 repo:
  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/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_8.1.1.33-1+cuda11.2_amd64.deb && \
  apt-get install -y ./libcudnn8-dev_8.1.1.33-1+cuda11.2_amd64.deb
fi

if [ $? -eq 0 ]; then
  #reboot system
  #systemctl reboot
  echo "\n"
  echo "Setup successfully finished"
  echo "Type \"systemctl reboot\" to reboot system"
fi