#!/bin/bash
######################################################################
#                           JURECA
#-------------------------------------------------------------------
# using SLURM
######################################################################

function UJS_Submit
{
	echo "Cluster: JURECA. Scheduler: SLURM."
	
	# check consistency
	if [ $((npe%nppn)) -ne 0 ]; then
		echo "npe=$npe is not divisible by nppn=$nppn"
		exit
	fi
	
	if [ $nnodes -gt 128 ]; then
		echo "Jureca does not provide more than 128 nodes for one job."
		return
	fi	
	
	## interactive (devel) jobs (not tested)
	# Once an allocation has been made, the salloc command will start a bash
	# on the login node where the submission was done. After a successful
	# allocation the users can execute srun from that shell and they can
	# spawn interactively their applications. The interactive session is
	# terminated by exiting the shell.
	# In order to obtain a shell on the first allocated compute nodes
	# (like command “msub -I“ from Moab), the users can start a remote shell
	# from within the current session and connect it to a pseudo terminal
	# (pty) using the srun command with a shell as an argument. For example:
	#     srun --cpu_bind=none --nodes=2 --pty /bin/bash
	# After gaining access to the remote shell it is possible to run srun
	# again from that remote shell in order to execute interactively
	# applications without any delays (no scheduling delays since the
	# allocation has already been granted).
	# For more details, cf.:
	# "JURECA - User's Manual for the Batch System - Slurm"
	if [ $interactive == true ]; then
		
		if [ $walltime == "unlimited" ]; then
			walltime=00:30:00
		fi
		
		# todo: check walltime is <= 2h
		
		if [ $nnodes -gt 8 ]; then
			echo "ERROR: The maximum number of nodes for interactive jobs is 8."
			return
		fi
		
		commandline="salloc -N $nnodes -n $npe --partition=devel --job-name=$jobname --time=$walltime"
		echo " command:      "$commandline >> info.txt	
	
		echo "[[ current cluster allocation"
		squeue
		echo "]] current cluster allocation"
		
		if [ $test == true ]; then
			echo "ONLY testing - NOT executing."
			echo "Submit/Start: $commandline"
			return
		fi
				
		echo "Start: $commandline"
		
		$commandline | tee $outdir/job.output
		return=$?	
		if [ ! $return == 0 ]; then
			echo "ERROR: srun returned $return. Job has NOT been started."
			exit
		fi
		
	## BATCH jobs
	else
		
		# walltime handling
		if [ $walltime == "unlimited" ]; then
			walltime=00:30:00
		fi
		
		echo "Create: $outdir/job.sh"
		
		#MYLDLPATH=/bgsys/drivers/ppcfloor/comm/lib/
		
		# mail notification handling
		jurecaNotification="NONE"
		if [ $mail == true ]; then 
			if [ -z "$UGSUBMIT_EMAIL" ]; then
				echo "please set UGSUBMIT_EMAIL or specify email with -email. Aborting."
				exit
			fi
		
			if [ $mailStart == true ]; then
				jurecaNotification="BEGIN"
			fi
			if [ $mailEnd == true ]; then
				if [ $jurecaNotification == "NONE" ]; then
					jurecaNotification="END"
				else
					jurecaNotification="ALL"
				fi
			fi
			if [ $mailError == true ]; then
				if [ $jurecaNotification == "NONE" ]; then
					jurecaNotification="FAIL"
				else
					jurecaNotification="ALL"
				fi
			fi	 
		fi
		
		# generic resources
		if [ -z "$jureca_mem" ]; then
			jureca_gres="mem128"
		elif [ "$jureca_mem" == "128" -o "$jureca_mem" == "256" -o "$jureca_mem" == "512" ]; then
			jureca_gres="mem${jureca_mem}"
		else
			echo "Jureca-memory parameter not set correctly (\"$jureca_mem\"). Valid values are: 128, 256, 512."
			return
		fi
		
		# partition
		if [ -z "$jureca_part" ]; then
			jureca_part="batch"
		fi
		if [ "$jureca_part" == "devel" ]; then 
			if [ ! "$jureca_gres" == "mem128" ]; then
				echo "The Jureca devel partition can only be used with 128G memory nodes."
				return
			fi
		elif [ "$jureca_part" == "batch" ]; then
			if [ ! "$jureca_gres" == "mem128" ]; then
				echo "The Jureca batch partition can only be used with 128G memory nodes."
				return
			fi
		elif [ "$jureca_part" == "mem256" ]; then 
			if [ ! "$jureca_gres" == "mem256" ]; then
				echo "The Jureca mem256 partition can only be used with 256G nodes. Set -Jureca-memory 256."
				return
			fi
		elif [ "$jureca_part" == "mem512" ]; then 
			if [ ! "$jureca_gres" == "mem512" ]; then
				echo "The Jureca mem512 partition can only be used with 512G nodes. Set -Jureca-memory 512."
				return
			fi
		else
			echo "Jureca-partition parameter not set correctly. Valid values are: devel, batch, mem512."
			return
		fi
		
		
		# write job script		
		cat > job.sh << EOF
#!/bin/bash
#SBATCH --job-name=$jobname
#SBATCH --nodes=$nnodes
#SBATCH --ntasks=$npe
#SBATCH --time=$walltime
#SBATCH --partition=$jureca_part
#SBATCH --gres=$jureca_gres
#SBATCH --error=job.error
#SBATCH --output=job.output
#SBATCH --mail-type=$jurecaNotification
#SBATCH --mail-user=$UGSUBMIT_EMAIL

$profilePrefix
srun -N $nnodes -n $npe -p $jureca_part --gres=$jureca_gres $executable $args
EOF

		# execute command (or only print it in test case)
		commandline="sbatch -N $nnodes -n $npe job.sh"
		echo " command:      "$commandline >> info.txt
	
		if [ $test == true ]; then
			echo "ONLY testing - NOT executing."
			echo "Submit/Start: $commandline"
			return
		fi
	
		echo "Submit: $commandline"
	 	commlineoutput=$($commandline)
	 	echo "$commlineoutput"
		jobid=$(echo $commlineoutput | sed 's/.*[^0-9]\([0-9]\+\)[^0-9]*$/\1/')
	fi
}



function UJS_GetOptions
{
	nppnmax=24
	pemax=3072	
}

function UJS_Info
{
	echo "UGSUBMIT Info for JURECA:"
	if [ ! -z $1 ] && [ $1 == "all" ]; then
		#echo "squeue -o \"%.7i %6C %6D %.20j %.10u %.8T %.10M %.10l %.6D %19R\""
		squeue -o "%.7i %6C %6D %.20j %.10u %.8T %.10M %.10l %.6D %19R"
	else
		#echo "squeue -u $USER -o \"%.7i %6C %6D %.20j %.10u %.8T %.10M %.10l %.6D %19R\""
		squeue -u $USER -o "%.7i %6C %6D %.20j %.10u %.8T %.10M %.10l %.6D %19R"
	fi	
}

function UJS_Cancel
{
	echo "Using SLURM on JURECA"
	if [ ! -z $1 ] && [ $1 == "all" ]; then
		echo "your jobs:"
		squeue -u $USER -o "%.7i %6C %6D %.50j %.10u %.8T %.10M %.10l %.6D %19R"
		echo " "
		read -p "Are you sure you want to cancel all your jobs (yes=j,J,y or Y) ? " -n 1 -r
		echo " "		
		if [[ $REPLY =~ ^[JjYy]$ ]]
		then			
		    scancel --user=$USER
		fi
		
	else
		scancel $1
	fi
}