#!/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 running list of jobs on cluster"
   echo ""
   echo "  Loops through an input job_txt (txt with command on each line)"
   echo ""
   echo "--------------------------------------------------------------------------"
   echo "Usage: $cmdname  <job_list>  <optional flags> "
   echo "--------------------------------------------------------------------------"
   echo ""
   echo "optional flags:"
   echo ""
   echo " -t : test-mode, don/'t actually submit any jobs"
   echo " -e : preserve environment variables (--export=ALL to sbatch)"
   echo " -A <account> : account to use for allocation (default: $CC_COMPUTE_ALLOC)"
   echo " -E <env_script>:  script to source for env setup"
   echo ""
   usage_job_templates
   usage_job_depends

 }


#get options in front 
while getopts "J" options; do
 case $options in
    J ) print_job_templates
	exit 1;;
    * ) usage
	exit 1;;
 esac
done

shift $((OPTIND-1))

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


output_dir=`realpath $PWD`

joblist=$1
shift 1
job_template=Regular
depends=""
testmode=0
keepenv=0
cc_account=$CC_COMPUTE_ALLOC
envscript=

while getopts "Jj:td:eA:E:" 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;; 
    e ) echo "	Preserve environment (--export=ALL)" >&2
	keepenv=1;;
    A ) echo "	Using allocation account: $OPTARG" >&2
	    cc_account=$OPTARG;;
    E ) echo "	Sourcing $OPTARG in each job" >&2
	envscript=$OPTARG;;
   * ) usage
	exit 1;;
 esac
done


#check subjlist validity and get # of subjects
if [ -f $joblist ]
then
	N=`cat  $joblist | wc -l`
else
 	echo "joblist does not exist: $joblist" >&2
	exit 1
fi

indices=1-$N

job_dir=$output_dir/jobs
mkdir -p $job_dir

jobname=${joblist##*/} #strip leading path
job=$job_dir/$jobname.$RANDOM.job

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

echo "#SBATCH --job-name=$jobname" >> $job
echo "#SBATCH --account=$cc_account" >> $job

#submit with array, using subj and subjlist, indexed by line number (1-N)
echo "#SBATCH --output=$job_dir/${jobname}.%A_%a.out" >> $job
echo "#SBATCH --array=$indices" >> $job

if [ -n "$envscript" ]
then
	echo "source $envscript" >> $job
fi

echo "cmd=\`head -n \$SLURM_ARRAY_TASK_ID `realpath $joblist`  | tail -n 1\`" >> $job
echo cd `pwd` >> $job
echo "echo CMD: \$cmd " >>$job
echo "echo START_TIME: \`date\`" >>$job
echo "export SCRATCH_DIR=/scratch/${USER}/\${SLURM_ARRAY_JOB_ID}_\${SLURM_ARRAY_TASK_ID}" >> $job
echo "mkdir -p \$SCRATCH_DIR" >> $job
echo "\$cmd" >> $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
	
if [ "$keepenv" == 1 ]
then
	keepopt="--export=ALL"
else
	keepopt=""
fi

if [ "$testmode" == 0 ]
then
message=`sbatch $keepopt --dependency=$depends $job`

# 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/$scriptname.$jobid.job

#move instance of this job script for provenance
mv  $job $exec_job

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

echo $jobid

exit 0
