#!/usr/bin/env bash
#
# Install project dependencies and requirements to run experiments on cluster.
# Commands can be run either in login node or in compute node.
# OS supported: Ubuntu 20.04.3 LTS (Focal Fossa)

set -o errexit
set -o pipefail
set -o nounset

# Variable should not be read-only because setup script already defines it and would cause conflict.
CONDA_ENV_NAME='masterarbeit_env'
readonly CUDA_VERSION='11.3'
readonly PYTHON_VERSION='3.8'

init_bashrc() {
  local -r src_bashrc='/home/rodrigo/bashrc'
  local -r dest_bashrc='/home/rodrigo/.bashrc'

  echo "Init bashrc to path $dest_bashrc on cluster"

  # Check if bashrc was already copied to cluster
  if [ ! -f "$src_bashrc" ]; then
    echo "ERROR: File $src_bashrc does not exist."
    echo "Please run script/transfer-build-to-cluster first to transfer bashrc to cluster."
    exit 1
  fi

  # Check if .bashrc already exists
  if [ -f "$dest_bashrc" ]; then
    echo "Back up existing .bashrc under $dest_bashrc"
    mv "$dest_bashrc" "$dest_bashrc".bak
  fi

  echo "Install .bashrc from $src_bashrc"
  ln -s "$src_bashrc" "$dest_bashrc"
}

install_miniconda() {
  local -r arch='x86_64'
  local -r miniconda_path="${HOME}/miniconda3"

  if [ -d "$miniconda_path" ]; then
    echo "Miniconda directory under $miniconda_path detected. Skipping installation."
    return
  fi

  echo "Install Miniconda"

  echo "Downloading Miniconda's installation script for arch $arch"
  wget "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${arch}.sh"

  echo "Installing Miniconda for arch $arch"
  bash "Miniconda3-latest-Linux-${arch}.sh"
}

source_setup() {
  echo "Source setup script to load conda environment, if installed"
  # Disable shellcheck warning about using source with a variable.
  # shellcheck disable=SC1091
  source /home/rodrigo/experiments/cluster/script/setup
}

update_conda() {
  echo 'Updating conda'
  # Command is shown when running a conda command in an outdated environment.
  conda update --name base --channel defaults conda
}

install_pytorch() {
  if conda env list | grep -q "$CONDA_ENV_NAME"; then
    echo "Conda environment $CONDA_ENV_NAME detected. Skipping installation."
    return
  fi

  echo 'Install PyTorch'

  echo "Creating new conda environment called '$CONDA_ENV_NAME'"
  conda create --name "$CONDA_ENV_NAME" python="$PYTHON_VERSION" --yes

  echo 'Activating created conda environment'
  conda activate "$CONDA_ENV_NAME"

  echo "Installing PyTorch with CUDA ($CUDA_VERSION) support"
  # Source: https://pytorch.org/get-started/locally/
  conda install pytorch torchvision torchaudio cudatoolkit="$CUDA_VERSION" --channel pytorch --yes
}

main() {
  echo "Bootstrap cluster node"
  init_bashrc
  install_miniconda
  source_setup
  update_conda
  install_pytorch
}

main
