#!/bin/bash
#
# Submit job for Quantum Espresso PWSCF program
# Created: 20160708
# Wirawan Purwanto
#

if [ -n "$PWRUN_DEBUG_SCRIPT" ]; then
  Debug=$PWRUN_DEBUG_SCRIPT
fi

: ${PWRUN_TRACE_SCRIPT:=0}
if [ "$PWRUN_TRACE_SCRIPT" = 1 ]; then
  set -x
  export
fi

set -e
shopt -s extglob

# Default settings -- be careful changing these.
: ${PWRUN_COMPUTE_HOST:=batch}
: ${PWRUN_BATCHSYS_JOBNAME:=pwscf-5.3}
QE_ROOT=/cm/shared/apps/qe/5.3-gcc
: ${Debug:=}
# Some cluster defaults (may change over time)
# Reference: /etc/profile.d/modules.sh
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
# - LOADEDMODULES

RestoreModuleEnv () {
  local LOADEDMODULES_SAVE
  # : ${MODULESHOME:=$MODULESHOME_DEFAULT}
  # : ${MODULEPATH:=$MODULEPATH_DEFAULT}
  MODULESHOME=$MODULESHOME_DEFAULT
  MODULEPATH=$MODULEPATH_DEFAULT
  export MODULESHOME
  export MODULEPATH
  if [[ ":${MODULEPATH}:" != *":${MODULEPATH_DEFAULT}:"* ]]; then
    MODULEPATH=${MODULEPATH}:${MODULEPATH_DEFAULT}
  fi

  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 "Loaded modules after module restore: $LOADEDMODULES"
    module list
  fi
}

InitModuleEnv () {
  # Only do the initialization part of the module so "module" command
  # is available for the script
  local LOADEDMODULES_SAVE
  : ${MODULESHOME:=$MODULESHOME_DEFAULT}
  : ${MODULEPATH:=$MODULEPATH_DEFAULT}
  export MODULESHOME
  export MODULEPATH

  source "$MODULESHOME/init/bash"
}

module_loaded () {
  # Usage: module_loaded  <module_pattern>
  # Returns code zero if the said module has been loaded, or nonzero if not.
  case ":$LOADEDMODULES:" in
    (*:$MOD:*)
      return 0
      ;;
    (*)
      return 1
      ;;
  esac
}


function path_search()
# Usage: path_search file path:path:path:...  [testop]
{
  local IFS=:
  local p d
  local testop="${3:-"-f"}"
  if [[ "$1" == /* ]]; then
    if test $testop "$1"; then
      echo "$1"
      return 0
    else
      return 1
    fi
  fi
  for d in $2; do
    p=$d/$1
    if test $testop "$p"; then
      echo "$p"
      return 0
    fi
  done
  return 1
}


Help () {
  echo "Usage: $MYSELF [number_of_processors] [inputfile] [outputfile]"
  echo 
  echo "Example: $MYSELF 4 MgO.in MgO.out"
}


Preamble () {
  MYSELF_PATH=$(path_search "$0" "$PATH")
  MYSELF=$(basename "$MYSELF_PATH")

  if ((Debug >= 10)); then
    echo "Env-vars:"
    export
  fi
}


ParseArgs () {
# FIXME: This is very primitive but accomodates what the original
# qe-q-5.3 provides.

  PWRUN_OUTPUT=
  if [ $# -lt 2 ]; then
    Help >&2
    exit 1
  fi
  PWRUN_NUM_TASKS=$1
  PWRUN_INPUT=$2

  if [ ! -f "$PWRUN_INPUT" ]; then
    echo "Error: input file $PWRUN_INPUT cannot be found" >&2
    exit 2
  fi

  if [ -z "$3" ]; then
    PWRUN_OUTPUT="${PWRUN_INPUT%.in}.out"
  else
    PWRUN_OUTPUT=$3
  fi

  if ((Debug > 0)); then
    echo "PWRUN_INPUT=$PWRUN_INPUT"
    echo "PWRUN_OUTPUT=$PWRUN_OUTPUT"
  fi

  if [ $# -gt 3 ]; then
    shift 3
    PWRUN_BATCHSYS_OPTS=( "$@" )
    if ((Debug > 0)); then
      echo "Added extra flags to batch system: ${PWRUN_BATCHSYS_OPTS[*]}"
    fi
  fi
}

SubmitBatchJob () {
  export PWRUN_COMPUTE_HOST=computenode
  local VARS
  VARS=LOADEDMODULES,MODULESHOME,MODULEPATH
  VARS=$VARS,PWRUN_COMPUTE_HOST,PWRUN_INPUT,PWRUN_OUTPUT,PWRUN_NUM_TASKS
  VARS=$VARS,Debug,PWRUN_TRACE_SCRIPT
  export ${VARS//,/" "}
  if ((Debug > 0)); then
    set -x
  fi
  qsub -cwd -S /bin/bash -v $VARS -pe qe "$PWRUN_NUM_TASKS" -j yes \
     "${PWRUN_BATCHSYS_OPTS[@]}" \
     "$MYSELF_PATH"
}

RunBatchJob () {
  RestoreModuleEnv

  # Provides some fallback modules if they are not loaded
  if ! module_loaded "gcc/*"; then
    if ((Debug > 0)); then
      echo "Note: adding fallback module gcc/5.3.0"
    fi
    module load gcc/5.3.0
  fi
  if ! module_loaded "openmpi/gcc/64/*"; then
    if ((Debug > 0)); then
      echo "Note: adding fallback module openmpi/gcc/64/1.10.2"
    fi
    module load openmpi/gcc/64/1.10.2
  fi

  # TEMPORARY FIX
  export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/cm/shared/apps/fftw/fftw-3.3.4-shared/lib

  QE_PROG=$QE_ROOT/pw.x

  if ((Debug > 0)); then
    echo "PWRUN_INPUT=$PWRUN_INPUT"
    echo "PWRUN_OUTPUT=$PWRUN_OUTPUT"

    if ((Debug > 1)); then
      echo "Loaded-modules::"
      module list >&1 || true
      echo ""
      echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"
      echo "QE_PROG=$QE_PROG"
      ldd $QE_PROG || true
    fi
    set -x
  fi

  mpirun -np "$PWRUN_NUM_TASKS"  "$QE_PROG" -i "$PWRUN_INPUT" > "$PWRUN_OUTPUT"
  exit $?
}



Main() {
  Preamble

  case "$PWRUN_COMPUTE_HOST" in
  (computenode)
    RunBatchJob
    ;;

  (batch)
    ParseArgs "$@"
    SubmitBatchJob    
    ;;
  esac
}

Main "$@"
exit $?

eval `/cm/local/apps/environment-modules/3.2.10/bin/modulecmd bash load /cm/shared/modulefiles/gcc/5.3.0`
eval `/cm/local/apps/environment-modules/3.2.10/bin/modulecmd bash load /cm/shared/modulefiles/openmpi/gcc/64/1.10.2`

name=`basename $0`

if [ $# -ne 3 ]; then 
        echo "Usage: $name [number of processors {1 - 256}] [inputfile] [outputfile]"
	echo 
	echo "Example: $name 4 data.in results.out"
	exit -1
fi


qrsh -V -cwd -N QE-SUB -pe qe $1 mpirun -np $1 /cm/shared/apps/qe/5.3-gcc/pw.x -i $2 > $3 &