#!/bin/bash

NB_RUNS="5"
KCAT_CAN_TOMO="1.0"
STIM_FREQ="10"
RUN_NAME="Vesicle_cycle_${STIM_FREQ}Hz"
SLICE_ENDT="10"
ENDT="50"

SLURM_PARAMS="--partition=compute --job-name=${RUN_NAME} --output=${RUN_NAME}_%A.out --nodes 1 --ntasks=128 --exclusive --mem-per-cpu=3700m --time=4-0:0:0"

SCRIPT_PARAMS=". meshes/pyr_axon_2021 logs/ data ${RUN_NAME} \${SLURM_JOB_ID} __TASK_ID__ ${KCAT_CAN_TOMO} ${STIM_FREQ} __ENDT__ false"
SCRIPT_PARAMS_INIT="checkpoint checkpoints/${RUN_NAME}___TASK_ID_____ENDT__s"
SCRIPT_PARAMS_RESTORE="both checkpoints/${RUN_NAME}___TASK_ID_____PREVT__s checkpoints/${RUN_NAME}___TASK_ID_____ENDT__s"

SLURM_SCRIPT_BASE="
#!/bin/bash

module load steps/STEPS_5_0_3

export UCX_MM_RX_BUFS_GROW=2048
srun --mpi=pmix python -u VESICLE_CYCLE_MODEL_MPI_API2.py"
SLURM_SCRIPT_INIT="${SLURM_SCRIPT_BASE} ${SCRIPT_PARAMS} ${SCRIPT_PARAMS_INIT}"
SLURM_SCRIPT_RESTORE="${SLURM_SCRIPT_BASE} ${SCRIPT_PARAMS} ${SCRIPT_PARAMS_RESTORE}"

for tid in $(seq 1 ${NB_RUNS})
do
	TIME=${SLICE_ENDT}

	# Try to find previously recorded checkpoints
	PREVT=0
	while [ -f checkpoints/${RUN_NAME}_${tid}_${TIME}s_0 ]
	do
		PREVT=$TIME
		if [ $(echo "$TIME + $SLICE_ENDT < $ENDT" | bc) -gt 0 ]
		then
			TIME=$(echo "${TIME} + ${SLICE_ENDT}" | bc)
		else
			TIME=$ENDT
		fi
	done

	# If no checkpoints were found, submit the first slice without restoring
	if [ ${PREVT} -eq 0 ]
	then
		SLURM_SCRIPT=${SLURM_SCRIPT_INIT//__TASK_ID__/${tid}}
		SLURM_SCRIPT=${SLURM_SCRIPT//__ENDT__/${TIME}}

		jobstr=$(sbatch ${SLURM_PARAMS} --wrap "${SLURM_SCRIPT}")
		jobid=${jobstr##* } 

		PREVT=$SLICE_ENDT
		TIME=$(echo "${TIME} + ${SLICE_ENDT}" | bc)

		echo "Submitted job ${jobid} for task ${tid} and slice 0s to ${TIME}s."
	else
		jobid=""
	fi

	if [ ${TIME} -le ${ENDT} ]
	then
		# Submit the remaining slices, starting with the found checkpoint, if any
		while [ $(echo "$PREVT < $ENDT" | bc) -gt 0 ]
		do
			SLURM_SCRIPT=${SLURM_SCRIPT_RESTORE//__PREVT__/$PREVT}
			SLURM_SCRIPT=${SLURM_SCRIPT//__ENDT__/$TIME}
			SLURM_SCRIPT=${SLURM_SCRIPT//__TASK_ID__/$tid}
			if [ -z $jobid ]
			then
				SLURM_DEPS=""
			else
				SLURM_DEPS="--dependency=afterok:${jobid} --kill-on-invalid-dep=yes"
			fi

			jobstr=$(sbatch ${SLURM_PARAMS} ${SLURM_DEPS} --wrap "${SLURM_SCRIPT}")
			jobid=${jobstr##* } 
			echo "Submitted job ${jobid} for task ${tid} and slice ${PREVT}s to ${TIME}s."

			PREVT=$TIME
			if [ $(echo "$TIME + $SLICE_ENDT < $ENDT" | bc) -gt 0 ]
			then
				TIME=$(echo "${TIME} + ${SLICE_ENDT}" | bc)
			else
				TIME=$ENDT
			fi
		done
	else
		echo "Run {tid} already finished, nothing to submit."
	fi
done
