#!/bin/bash

# ChemShell SLURM batch script
# you.lu@stfc.ac.uk

# unload modules
for mod in ${CHEMSH_UNLOAD_MODS}
do
    module unload $mod
done

# restore
#module restore

# swap modules
IFS=$';'
for pair in ${CHEMSH_SWAP_MODS}
do
    IFS=$' ' read -a mods <<< $pair
    module swap ${mods[0]} ${mods[1]}
done

# load modules
IFS=$' '
for mod in ${CHEMSH_LOAD_MODS}
do
    module load $mod
done

# epcc-job-env unloads cray-python
module load cray-python

# debug
#module --quiet load valgrind4hpc/2.10.2

module list

echo "Node list is" $SLURM_NODELIST

# some things that might be useful
echo
echo SLURM parameters:
echo
echo SLURM_JOBID         = $SLURM_JOBID
echo SLURM_JOB_NODELIST  = $SLURM_JOB_NODELIST
echo SLURM_NNODES        = $SLURM_NNODES
echo SLURMTMPDIR         = $SLURMTMPDIR
echo SLURM_ARRAYID       = $SLURM_ARRAYID
echo SLURM_ARRAY_JOB_ID  = $SLURM_ARRAY_JOB_ID
echo SLURM_ARRAY_TASK_ID = $SLURM_ARRAY_TASK_ID
echo SLURM_CPUS_PER_TASK = $SLURM_CPUS_PER_TASK
echo SLURM_CPU_FREQ_REQ  = $SLURM_CPU_FREQ_REQ

NPROCS=`srun --nodes=${SLURM_NNODES} bash -c 'hostname' | wc -l`

echo
echo ChemShell runtime parameters:
echo
echo working directory = $SLURM_SUBMIT_DIR
echo PYTHONPATH        = $PYTHONPATH
echo NPROCS            = $NPROCS
echo NTHREADS          = $CHEMSH_NTHREADS
echo WRAPPER           = $CHEMSH_WRAPPER
echo EXEC              = $CHEMSH_EXEC
echo NWORKGROUPS       = $CHEMSH_NWORKGROUPS
echo INPUT             = $CHEMSH_INPUT
echo OUTPUT            = $CHEMSH_OUTPUT
echo Print level       = $CHEMSH_DEBUG
echo Modules to load   = $CHEMSH_LOAD_MODULES
echo Modules to unload = $CHEMSH_UNLOAD_MODULES
echo Modules to swap   = $CHEMSH_SWAP_MODULES

# TWK 17/7/2023 Test driver mode
if [[ "$CHEMSH_TEST_DRIVER" != "" ]]
then
    echo 
    echo Running the ChemShell test suite:
    echo 
    echo CHEMSH_TEST_DRIVER = $CHEMSH_TEST_DRIVER
    echo CHEMSH_TESTSETS = $CHEMSH_TESTSETS
    echo
    CHEMSH_TEST="${CHEMSH_TEST_DRIVER} --sets ${CHEMSH_TESTSETS} --exe "
else
    CHEMSH_TEST=""
fi

# use srun preferably to use the allocated resources in a job step

# OpenMPI
export OMP_NUM_THREADS=$CHEMSH_NTHREADS
# TWK 4/1/2024 Require one thread per physical core (see ARCHER2 documentation)
if [[ $CHEMSH_NTHREADS -gt 1 ]]
then
    export OMP_PLACES=cores
fi

# Propagate the cpus-per-task setting from script to srun commands
# as recommended in the ARCHER2 user documentation.
#    By default, Slurm does not propagate this setting from the sbatch
#    options to srun commands in the job script. If this is not done,
#    process/thread pinning may be incorrect leading to poor performance
export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK

export ATP_ENABLED=1
# GA/ARMCI
export ARMCI_DEFAULT_SHMMAX_UBOUND=131072

# debug
#valgrind4hpc --tool=memcheck --valgrind-args="--track-origins=yes --leak-check=full --show-leak-kinds=all" --outputfile=_valgrind.log --launcher-args="-n 128" srun -- $CHEMSH_EXEC -nwg $CHEMSH_NWORKGROUPS -npmm $CHEMSH_NPROCS_MM -dbg $CHEMSH_DEBUG $CHEMSH_INPUT > $CHEMSH_OUTPUT

# srun verbose mode
#ulimit -c unlimited
#srun --slurmd-debug=verbose -n $CHEMSH_NPROCS $CHEMSH_EXEC -nwg $CHEMSH_NWORKGROUPS -npmm $CHEMSH_NPROCS_MM -dbg $CHEMSH_DEBUG $CHEMSH_INPUT

# ATP debug mode
#ATP_ENABLED=1 srun -n $CHEMSH_NPROCS $CHEMSH_EXEC -nwg $CHEMSH_NWORKGROUPS -npmm $CHEMSH_NPROCS_MM -dbg $CHEMSH_DEBUG $CHEMSH_INPUT

echo MEM_ARCHER2 = $MEM_ARCHER2

echo Executing command:
if [[ $MEM_ARCHER2 -gt 0 ]]
# YL 21/12/2021: request a certain amount of memory for ChemShell and reserve the rest for system call using srun (e.g., for NAMD)
then
    ARCHER2_SRUN_OPTIONS="--oversubscribe --mem=${MEM_ARCHER2}M"
else
    # TWK 4/1/2024 Standard srun options recommended in ARCHER2 user documentation
    ARCHER2_SRUN_OPTIONS="--distribution=block:block --hint=nomultithread"
fi

echo $CHEMSH_TEST $CHEMSH_WRAPPER srun -n $CHEMSH_NPROCS $ARCHER2_SRUN_OPTIONS $CHEMSH_EXEC -nwg $CHEMSH_NWORKGROUPS -npmm $CHEMSH_NPROCS_MM -gm $CHEMSH_GA_MEMORY -dbg $CHEMSH_DEBUG $CHEMSH_INPUT \> $CHEMSH_OUTPUT
$CHEMSH_TEST srun -n $CHEMSH_NPROCS $ARCHER2_SRUN_OPTIONS $CHEMSH_EXEC -nwg $CHEMSH_NWORKGROUPS -npmm $CHEMSH_NPROCS_MM -gm $CHEMSH_GA_MEMORY -dbg $CHEMSH_DEBUG $CHEMSH_INPUT > $CHEMSH_OUTPUT

# debug mode
#export PYTHONPATH=/work/c01/c01/ylu/chemsh/archer2
#export CHEMSH_ARCH=cray
# NB: because MPI is not initialised code involving MPI will fail in this mode (including base.connect)
#valgrind --tool=memcheck --leak-check=full --show-leak-kinds=all --log-file=_valgrind.log python3 $CHEMSH_INPUT

