#!/bin/bash
#
# submit invokes picmc via SLURM
#
# Andreas Pflug 2017-01-20
# Based on rvmsge by Christoph Schwanke 2011-02-24
#
# Usage: rvmslurm [-help] [-picmc PICMCTasks] [-bem BEMTasks] [ -time duration_hours ] [ -mem memory_by_core_in_MB ] [ -partition slurm_partition ] <batchfile> 

# ======================
# Editable value section
# ======================

picmcpath=$PICMCPATH                           # PICMC installation path
partition="v3"                                 # Default partition
memory=1300                                    # Default memory in MB per CPU
walltime="336:00:00"		               # Default job duration in hhh:mm:ss (336 hours -> 14 days), 
				               # specify -time <nn> of command line for different job duration
MPIEXEC=mpiexec                                # Name of mpiexec executable
MPIPATH=/opt/openmpi-2.1.0                     # Full path to OpenMPI installation
#MAILRECIPIENT=user@institution.com            # Uncomment this to get status emails from SLURM 
                                               # (Needs a running mail server installation such as postfix)

QSUB=sbatch

# =============================
# End of editable value section
# =============================

numPICMCTasks=0
numBEMTasks=0
batchfile=""
fileExist=false
after=""

help=false

HELP="Help: Invokes rvmmpi via SLURM batch scheduling system"
USAGE="Usage: rvmslurm [-help] [-picmc PICMCTasks] [-bem BEMTasks] [-partition <partition>] [-N <number of nodes>] [-time <job duration in hh:mm:ss> ] [-mem <memory_by_core_in_MB>] [-partition <slurm_partition> ] <parameter file>"

while [ $# -gt 0 ]

do case "$1" in

	-help) help=true
	shift;;

	-mem) 
	    shift
	    if [[ $1 != [0-9]* ]] ; then
		echo "ERROR: $1 is not a valid number of memory"
		exit 1
	    fi
	    memory=$1
	    shift;;

	-after)
	    shift
	    after=$1
	    shift;;

	-partition)
	    shift
	    partition=$1
	    shift;;

	-N)
	    shift
	    NODES=$1
	    shift;;

	-picmc) shift
		if [[ $1 != [0-9]* ]] ; then
			echo "ERROR: $1 is not a valid number of PICMC Tasks"
			exit 1
		fi
		if [ $# -eq 0 ]
			then echo "no PICMC Tasks specified"
		else numPICMCTasks=$1
		fi
		shift;;

 	-time) shift
		if [[ $1 != [0-9]* ]] ; then
			echo "ERROR: $1 is not a valid number of hours"
			exit 1
		fi
		walltime=$1
		shift;;

	-bem) shift

		if [[ $1 != [0-9]* ]]
			then echo "ERROR: $1 is not a valid number of BEM Tasks"
			exit 1
		fi
		if [ $# -eq 0 ]
			then echo "no BEM Tasks specified"
		else numBEMTasks=$1
		fi
		shift;;

	*)	batchfile="$1"
		if [ $numPICMCTasks -gt 0 ]
			then if [ -f $batchfile ]
				then fileExist=true
			elif [ -f $batchfile".par" ]
				then fileExist=true
			fi
			if [ $fileExist = "false" ]
				then echo "ERROR: could not find a .par file"
				echo "$USAGE"
				exit 1
			fi
		fi
		if [ $numBEMTasks -gt 0 ]
			then fileExist=false
			if [ -f $batchfile ]
				then fileExist=true
			fi
			if [ -f $batchfile".par" ]
				then fileExist=true
			fi
		fi
		shift;;
	esac
done

if $help
	then echo "$HELP"
	echo "$USAGE"
	exit 1
fi

maxTasks=$(expr $numPICMCTasks + $numBEMTasks)

if [ $maxTasks -eq 0 ]
	then echo "ERROR: no tasks specified"
	echo "$USAGE"
	exit 1
fi

if [ $fileExist = "false" ]
	then echo "ERROR: no batchfile specified"
	echo "$USAGE"
	exit 1
fi

maxTasks=$(expr $maxTasks + 1)

jobdir=`pwd`;

slurmbatchfile="${batchfile}.slurm"

cat << EOF > ${slurmbatchfile}
#!/bin/bash
#SBATCH --partition=${partition}
#SBATCH --ntasks=${maxTasks}
#SBATCH --time=${walltime}
#SBATCH --mem-per-cpu=${memory}
#SBATCH -J ${batchfile}
#SBATCH -e "${batchfile}-err.txt"
#SBATCH -o "${batchfile}-out.txt"
EOF
if [ "${MAILRECIPIENT}" != "" ]; then
  echo "#SBATCH --mail-user=$MAILRECIPIENT" >> ${slurmbatchfile}
  echo "#SBATCH --mail-type=ALL" >> ${slurmbatchfile}
fi

if [ "${NODES}" != "" ]; then
  echo "#SBATCH -N ${NODES}" >> ${slurmbatchfile}
  echo "#SBATCH --ntasks-per-node=$((maxTasks/NODES+1))" >> ${slurmbatchfile}
fi

if [ "${after}" != "" ]; then
  echo "#SBATCH --dependency=afterok:${after}" >> ${slurmbatchfile}
fi

cat << EOF >> ${slurmbatchfile}
export PICMCPATH=$picmcpath
export LD_LIBRARY_PATH=$picmcpath/lib:/lib64:\$LD_LIBRARY_PATH
export PATH=$MPIPATH/bin:$picmcpath/bin:\$PATH
cd $jobdir
EOF

if [ $numBEMTasks -gt 0 ] ; then
	fileexists=false
	if [ -f $batchfile".bem" ]
		then initbfield="$picmcpath/sh/initbfield $batchfile"
		$initbfield
		fileexists=true
	fi

	if [ $fileexists = "false" ]
		then echo "ERROR: could not find magnetic field computation template $batchfile.bem"
		echo "Execute initbfield command to create it"
		exit 1
	fi

	fileexists=false
	if [ -f $batchfile ]
		then fileexists=true
	fi

	if [ -f $batchfile".rd" ]
		then bemfile=$batchfile".rd"
		fileexists=true
	fi
	
	if [ -f $batchfile"_bfield.rd" ]
		then bemfile=$batchfile"_bfield.rd"
		fileexists=true
	fi

cat << EOF >> ${slurmbatchfile}
${MPIEXEC} --bind-to none -x LD_LIBRARY_PATH -x PICMCPATH -np 1 \$PICMCPATH/bin/rvm -i \$PICMCPATH $jobdir/$bemfile : -np $numBEMTasks \$PICMCPATH/bin/bem
EOF

fi

if [ $numPICMCTasks -gt 0 ] ; then
cat << EOF >> ${slurmbatchfile}
${MPIPATH}/bin/${MPIEXEC} --prefix ${MPIPATH} --bind-to none -np 1 -mca mca_base_env_list "LD_LIBRARY_PATH=$LD_LIBRARY_PATH;PICMCPATH=$PICMCPATH" \$PICMCPATH/bin/rvm -i \$PICMCPATH -l scr/master.r ${batchfile}.par : -np $numPICMCTasks -mca mca_base_env_list "LD_LIBRARY_PATH=$LD_LIBRARY_PATH;PICMCPATH=$PICMCPATH" \$PICMCPATH/bin/picmc
EOF
fi

$QSUB $batchfile.slurm
