#!/bin/bash
######################################################################
#                           Shaheen II
#-------------------------------------------------------------------
# using SLURM
######################################################################

function UJS_Submit
{
	echo "Cluster: Shaheen II. Scheduler: SLURM."
	
	# check consistency
	if [ $((npe%nppn)) -ne 0 ]; then
		echo "npe=$npe is not divisible by nppn=$nppn"
		exit
	fi
	
	## interactive (devel) jobs: Not implemented in this script
	if [ $interactive == true ]; then
		
		echo "Interactive mode is not implemented in this script"
		exit
				
	## BATCH jobs
	else
		
		# walltime handling
		if [ $walltime == "unlimited" ]; then
			# Shaheen only permits 24 hours walltime in the workq partition
			walltime=23:59:00
		fi
		
		echo "Create: $outdir/job.sh"
		
		# mail notification handling
		shaheenNotification="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
				shaheenNotification="BEGIN"
			fi
			if [ $mailEnd = true ]; then
				if [ $shaheenNotification == "NONE" ]; then
					shaheenNotification="END"
				else
					shaheenNotification="ALL"
				fi
			fi
			if [ $mailError = true ]; then
				if [ $shaheenNotification == "NONE" ]; then
					shaheenNotification="FAIL"
				else
					shaheenNotification="ALL"
				fi
			fi	 
		fi
		
		# partition
		if [ -z "$partition" ]; then
			partition="workq"
		elif [ ! "$partition" = "72hours" ] && [ ! "$partition" = "debug" ]; then
			echo "Partition parameter not set correctly. Valid values are: workq, 72hours and debug."
			return
		fi
		
		# project for accounting
		if [ -z "$shaheen_project" ]; then
			echo "Shaheen project parameter must be set (for accounting)."
			return
		fi
		
		
		# write job script		
		cat > job.sh << EOF
#!/bin/bash
#SBATCH --account=$shaheen_project
#SBATCH --job-name=$jobname
#SBATCH --nodes=$nnodes
#SBATCH --time=$walltime
#SBATCH --partition=$partition
#SBATCH --error=job.error
#SBATCH --output=job.output
EOF

		if [ $mail = true ]; then 
		cat >> job.sh << EOF
#SBATCH --mail-type=$shaheenNotification
#SBATCH --mail-user=$UGSUBMIT_EMAIL
EOF
		fi

		cat >> job.sh << EOF
$profilePrefix
srun --ntasks $npe --tasks-per-node $nppn $executable $args
EOF

		# execute command (or only print it in test case)
		commandline="sbatch 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=32
	pemax=197568
}

function UJS_Info
{
	echo "UGSUBMIT Info for Shaheen:"
	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 Shaheen"
	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
}