#!/bin/bash

execpath=`dirname $0`
execpath=`realpath $execpath`


source ${NEUROGLIA_BASH_LIB}
if [ ! "$?" = 0 ]
then
	echo "Error initializing neuroglia-helpers, quitting $0"
	exit 1
fi


function usage {

   echo""
   echo "=========================================================================="
   echo "Interface for submitting a single cluster job with singularity"
   echo ""
   echo "--------------------------------------------------------------------------"
   echo "Usage: $cmdname <optional flags>  <command to run>"
   echo "--------------------------------------------------------------------------"
   echo ""
   echo "optional flags:"
   echo ""
   echo " -t : test-mode, don/'t actually submit any jobs"
   echo " -n : ninja-mode, delete job file and hide command in log file"
   echo " -r : run-mode, use singularity run instead of singularity exec"
   echo " -U <uri> : uri for singularity container to use (default: $NEUROGLIA_URI)"
   echo " -I <simg> : path to container to use (default: `shub-cache $NEUROGLIA_URI`)"
   echo " -N <name> : name the batch job (default: name is command with path stripped off)"
   echo " -A <account> : account to use for allocation (default: $CC_COMPUTE_ALLOC)"
   echo ""
   usage_job_templates
   usage_job_depends

}


#run a pipeline job on the cluster
#pipeline jobs take subjids as args
cmd=$0
cmdname=${cmd##*/}


if [ "$#" -lt 1 ]
then 
  usage
  exit 1
fi

job_template=Regular
depends=""
testmode=0
enable_wait=0
ninja=0
scmd="exec"
uri=$NEUROGLIA_URI
SINGULARITY_IMG=
output_dir=`realpath $PWD`
scriptname=
cc_account=$CC_COMPUTE_ALLOC

while getopts "Jj:d:tWnrN:I:A:U:" options; do
 case $options in
    J ) print_job_templates
	exit 1;;
    j ) echo "	Using job template: $OPTARG" >&2
	job_template=$OPTARG;;
    d ) echo "	Using dependencies: $OPTARG" >&2
	depends=$OPTARG;;
    t ) echo "	Using test-mode (no submit jobs)" >&2
	testmode=1;;
    W ) echo "  Wait until job completes" >&2
	enable_wait=1;;
    n ) echo "  Using ninja mode" >&2
	ninja=1;;
    r ) echo "  Run instead of exec" >&2
	scmd="run";;
    N ) echo "	Using job name: $OPTARG" >&2
	scriptname=$OPTARG;;
    I ) echo "	Using SINGULARITY_IMG: $OPTARG" >&2
	SINGULARITY_IMG=$OPTARG;;
    U ) echo "  Using singlarity image uri: $OPTARG" >&2
	uri=$OPTARG;;
    A ) echo "	Using allocation account: $OPTARG" >&2
	    cc_account=$OPTARG;;
    * ) usage
	exit 1;;
 esac
done


shift $((OPTIND-1))

run_cmd=$@

if [ ! -n "$scriptname" ]
then
scriptname=${run_cmd%%\ *} #remove all but first word
scriptname=${scriptname##*/} #remove up to leading slash
fi

echo $scriptname >&2

if [ ! -n "$SINGULARITY_IMG" ]
then
  SINGULARITY_IMG=`shub-cache $uri`
  if [ ! "$?" = 0 ]
  then
    echo "error loading $uri, exiting!"
    exit 
  fi
fi

job_dir=$output_dir/jobs
mkdir -p $job_dir

job=$job_dir/$scriptname.$RANDOM.job

execpath=`dirname $0`
execpath=`realpath $execpath`
cp ${NEUROGLIA_DIR}/job-templates/${job_template}.job $job  

echo "#SBATCH --job-name=$cmdname" >> $job
echo "#SBATCH --account=$cc_account" >> $job
echo "#SBATCH --output=$job_dir/${cmdname}.%A.out" >> $job
echo "echo SINGULARITY_IMG: $SINGULARITY_IMG" >> $job
if [ "$ninja" = 1 ]
then
echo "echo Ninja-mode - delete job file" >> $job
else
echo "echo CMD: $@" >>$job
fi
echo "echo START_TIME: \`date\`" >>$job
echo "export SCRATCH_DIR=/scratch/${USER}/\${SLURM_JOB_ID}" >> $job
echo "mkdir -p \$SCRATCH_DIR" >> $job
echo cd `pwd` >> $job
echo singularity $scmd $SINGULARITY_OPTS $SINGULARITY_IMG $@ >> $job
echo "RETURNVAL=\$?" >> $job
echo "rm -rf \$SCRATCH_DIR" >> $job
echo "echo RETURNVAL=\$RETURNVAL" >>$job
echo "echo END_TIME: \`date\`" >>$job
echo "exit \$RETURNVAL" >> $job
echo -n "		Queuing job, $depends ... " >&2


if [ "$testmode" = 0 ]
then

if [ "$enable_wait" == 1 ]
then
message=`sbatch --dependency=$depends -W $job`
else
message=`sbatch --dependency=$depends $job`
fi

# Extract job identifier from SLURM's message.
if ! echo ${message} | grep -q "[1-9][0-9]*$"; then 
	echo "Job(s) submission failed." >&2
	echo ${message} >&2
	exit 1
else
	jobid=$(echo ${message} | grep -oh "[1-9][0-9]*$")
fi
echo "jobid=$jobid" >&2
exec_job=$job_dir/$cmdname.$jobid.job


if [ "$ninja" = 0 ]
then
#move instance of this job script for provenance
mv  $job $exec_job
fi

else
	jobid=$RANDOM
	echo "fake-jobid=$jobid  (test-mode: no jobs submitted)" >&2

fi


if [ "$ninja" = 1 ]
then
	#delete job
	rm -f $job
fi

echo $jobid

exit 0
