#
# DESCRIPTION gromacs-md-mdrun Molecular dynamics trajectory computation running GROMACS mdrun
#

# software_usage [optional]

# software_parse_cmd_options [optional]

set_software_environment () {
    JOBNAME=mds_mdrun_$(date +%s)
    [ -n "$MD_TPR" -a -f "$MD_TPR" ] || exit_error "Cannot find specified TPR file. Job submission aborted."
    if [ -n "$MD_CPT" ]; then
      if [ -f "$MD_CPT" ]; then 
        echo_g "GROMACS mdrun will be started from checkpoint (in accordance to params file)."
      else
        exit_error "Checkpoint file is defined in params but not found. Job submission aborted."
      fi
    fi
}

set_software_grid_environment () {
    GROMACS_RUNTIME=${GROMACS_RUNTIME:-GROMACS-5.0.4}
    GROMACS_VERSION=${GROMACS_RUNTIME##*-}
    if [ "$BUILD_RUNTIME" = "yes" ]; then
        next_idx_var INPUT_FILES_NAME_ INPUT_FILES_LOCATION_ \
                     "gromacs_install.sh" "${MDS_SOFT_INSTALLS_PATH%/}/GROMACS/gromacs_install.sh"
        next_idx_var INPUT_FILES_NAME_ INPUT_FILES_LOCATION_ \
                     "gromacs-${GROMACS_VERSION}.tar.gz" "ftp://ftp.gromacs.org/pub/gromacs/gromacs-${GROMACS_VERSION}.tar.gz"
    else
        next_idx_var RUNTIME_CONDITION_ RUNTIME_VALUE_ "=" "${GROMACS_RUNTIME}"
    fi
}

software_gen_params () {
cat <<END
###############################################################
#           GROMACS (MD) MDRUN SOFTWARE PARAMETERS            #
###############################################################

# GROMACS release (according to MolDynGrid runtime specification)
GROMACS_RUNTIME="GROMACS-5.0.4"

# Install GROMACS in runtime instead or requesting pre-installed one
#BUILD_RUNTIME=no (DO NOT ENABLE IN PRODUCTION YET)

# MDRUN topology input file
MD_TPR=md.tpr

# MDRUN options 
#  Output files will be generated automatically based on TPR file name.
#  This includes:
#      -o <name>.trr   -x <name>.xtc   -e <name>.edr    -c <name>.gro
#      -g <name>.log   -cpo <name>.cpt
#  You should specify only other options
MDRUN_OPTIONS="-v -nice 0 -dlb auto"

# MDRUN checkpoint file (specify to restart from checkpoint)
#MD_CPT=md.cpt

# Custom job name (redefine moldynsub autogenerated one)
#JOBNAME=my_mdrun_job

END
}

create_executable () {
cat << END
#!/bin/bash

# defines from moldynsub job params
BUILD_RUNTIME=${BUILD_RUNTIME}
GROMACS_VERSION=${GROMACS_VERSION}
MD_TPR=${MD_TPR}
MD_CPT=${MD_CPT}
MDRUN_OPTIONS="${MDRUN_OPTIONS}"

# CPUs topology (is required to correctly run under LRMS)
PPN=$PPN
NODES=$NODES
CPUS=$CPUS
[ -z "\$PPN" ] && PPN=\${CPUS:-1}
[ -z "\$NODES" ] && NODES=1


# Get base filename from TPR file
FBASE="\${MD_TPR%.tpr}"

# Install GROMACS in runtime if requested
if [ "\$BUILD_RUNTIME" = "yes" ]; then
  echo "INFO: Installing GROMACS as a part of job lifecycle"
  if [ -f "gromacs_install.sh" ]; then
    source gromacs_install.sh
  else
    echo "ERROR: There is no GROMACS install script found. Skipping."
  fi
else
  echo "INFO: Relying on pre-installed GROMACS software"
fi

# choose GROMACS mdrun flavor
if [ "\$NODES" != 1 ]; then
  echo "INFO: Computing on \$NODES nodes is requested. Using MPI-flavored GROMACS build."
  if type mdrun_mpi >/dev/null 2>&1; then
    MDRUN=mdrun_mpi
  else
    echo "ERROR: Cannot find mdrun_mpi executable. Job execution aborted."
    exit 1
  fi
else
  echo "INFO: Single-node computing is requested. Using GROMACS SMP build without MPI."
  unset GMX_MPIRUN
  export OMP_NUM_THREADS=\$PPN
  if type mdrun_smp >/dev/null 2>&1; then
    MDRUN=mdrun_smp
  elif type mdrun >/dev/null 2>&1; then
    MDRUN=mdrun
  else
    echo "FATAL: There is no GROMACS mdrun executable found. It seams GROMACS installation is broken."
    exit 1
  fi
fi

if [ ! -f "\${MD_TPR}" ]; then
  echo "FATAL: Cannot find defined TPR file (\${MD_TPR}) in current working directory. It seams submission has broken."
  exit 1
fi

if [ -n "\${MD_CPT}" ]; then 
  if [ -f "\${MD_CPT}" ]; then
    echo "INFO: Starting MD simulation from provided checkpoint"
    mv "\${MD_CPT}" mdrun_in_cpt.cpt
    MD_IN_CPT="-cpi mdrun_in_cpt.cpt"
  fi
fi

# Invoke MDRUN

\${GMX_MPIRUN} \${GMX_PREFIX} \${MDRUN} -s "\${MD_TPR}" -o "\${FBASE}.trr" -x "\${FBASE}.xtc" -e "\${FBASE}.edr" -c "\${FBASE}.gro" -g "\${FBASE}.log" -cpo "\${FBASE}.cpt" \${MD_IN_CPT} \${MDRUN_OPTIONS} 2>&1

END
}

create_input_files_list () {
    # only TPR and optionally checkpoint file are input files
    next_idx_var INPUT_FILES_NAME_ INPUT_FILES_LOCATION_ "$MD_TPR" ""
    [ -n "$MD_CPT" ] && next_idx_var INPUT_FILES_NAME_ INPUT_FILES_LOCATION_ "$MD_CPT" ""
    return 0
}

create_output_files_list () {
    local stageout_uri=$1
    # output files are set of script pre-defined files (TODO: handling of really optional files specified in extra arguments)
    local FBASE="${MD_TPR%.tpr}"
    for ofile in "${FBASE}.trr" "${FBASE}.xtc" "${FBASE}.edr" "${FBASE}.gro" "${FBASE}.log" "${FBASE}.cpt" ; do
      next_idx_var OUTPUT_FILES_NAME_ OUTPUT_FILES_LOCATION_ "$ofile" "${stageout_uri:+$stageout_uri/$ofile}"
    done
    return 0
}