# -*- bash -*-
#
# BASH support scriplet for environment modules.
#
# Simply "source" this script with "init" argument,
# then you can use "module" command from the batch script.
#
# Example (from within your batch script):
#
#     source ~/local/bash-module-env.sh init
#     module load gcc
#
# Optional: if you want to restore the modules listed in $LOADEDMODULES
# (which will need to be passed on to the script via "-v LOADEDMODULES" argument
# when qsub-ing, then use "restore" argment instead:
#
# Example (from within your batch script):
#
#     source ~/local/bash-module-env.sh restore
#

# Some cluster defaults (may change over time)
# Reference: /etc/profile.d/modules.sh
# These are the *old* module environment, pre-2016 upgrade
MODULESHOME_DEFAULT=/cm/local/apps/environment-modules/3.2.10/Modules/3.2.10
MODULEPATH_DEFAULT=/cm/local/modulefiles:/cm/shared/modulefiles:/cm/shared/compilers

# RestoreModuleEnv restores module environment at batch-job runtime to be
# exactly the same as the environment at the time of job submission.
# It requires the following variables to be set:
# - MODULESHOME (optional; default provided above)
# - MODULEPATH
# - LAODEDMODULES

RestoreModuleEnv () {
  local LOADEDMODULES_SAVE

  # Update 20170314: now Turing has support for old-style module even in
  # batch job. We will leverage that if we find that!
  # For proper operation, both LOADEDMODULES and _LMFILES_ must be exported
  # into the running environment.
  if [ -n "${MODULESHOME}" \
       -a -n "${MODULEPATH}" \
       -a "$(type -t module)" = "function" ]; then
    return 0
  fi
  case "${MODULESHOME}" in
  */lmod*)
    # FORCE reloading OLD modules for now.
    MODULESHOME=$MODULESHOME_DEFAULT
    ;;
  *)
    : ${MODULESHOME:=$MODULESHOME_DEFAULT}
    ;;
  esac
  _AppendPaths MODULEPATH "$MODULEPATH_DEFAULT"
  export MODULESHOME
  export MODULEPATH

  LOADEDMODULES_SAVE=$LOADEDMODULES
  unset LOADEDMODULES

  source "$MODULESHOME/init/bash"

  if [ -n "$LOADEDMODULES_SAVE" ]; then
    module load ${LOADEDMODULES_SAVE//:/" "}
  fi

  if [ -n "$Debug" -o -n "$DEBUG_MODULES" ]; then
    echo "New loaded modules: $LOADEDMODULES"
    module list
  fi
}

InitModuleEnv () {
  # Only do the initialization part of the module so "module" command
  # is available for the script
  if [ -n "${MODULESHOME}" \
       -a -n "${MODULEPATH}" \
       -a "$(type -t module)" = "function" ]; then
    return 0
  fi
  case "${MODULESHOME}" in
  */lmod*)
    # FORCE reloading OLD modules for now.
    MODULESHOME=$MODULESHOME_DEFAULT
    ;;
  *)
    : ${MODULESHOME:=$MODULESHOME_DEFAULT}
    ;;
  esac
  _AppendPaths MODULEPATH "$MODULEPATH_DEFAULT"
  #: ${MODULEPATH:=$MODULEPATH_DEFAULT}
  export MODULESHOME
  export MODULEPATH

  source "$MODULESHOME/init/bash"
}

_AppendPath () {
  # Usage: AppendPath VARNAME A_PATH
  # Appends a path element if it doesnt exist in current path
  local VAR VAL P
  VAR=$1
  P=$2
  VAL=$(eval "echo \"\$$VAR\"") # gets the current value of variable ${!VAR}
  case ":${VAL}:" in
  *:"${P}":*)
    ;;
  *)
    VAL="${VAL}:${P}"
    eval "$VAR=\$VAL"
    ;;
  esac
}

_AppendPaths () {
  # Usage: AppendPath VARNAME PATHLIST
  # Appends one or more elements from PATHLIST for the elements that don't
  # exist in current path
  local VAR PATH1 PATHLIST
  VAR=$1
  PATHLIST=$2
  for PATH1 in ${PATHLIST//:/" "}; do
    _AppendPath "$VAR" "$PATH1"
  done
}

# For convenience:
if [ "$1" = restore ]; then
  RestoreModuleEnv
elif [ "$1" = init ]; then
  InitModuleEnv
fi