#!/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 bids apps on the cluster with singularity"
   echo ""
   echo "  Can be used to run scripts that generally take command-line parameters as:"
   echo ""
   echo "Usage: $cmdname <bidsBatch options> <app_name> <bids_dir> <output_dir> <participant/group> <app options>"
   echo ""
   echo "Available apps:"
   print_app_list_names
   echo "  ( for list with url, use -B option )"
   echo ""
   echo "bidsBatch options:"
   echo ""
   echo " -s <subjid> : single-subject mode, run on a single subject (must be in participants.tsv) instead"
   echo " -S <subjlist> : subjlist mode, run on a subset of subjects defined by a simple text file (one id per line)"
   echo " -t : test-mode, don't actually submit any jobs"
   echo " -L <N> : local mode, run jobs locally with N parallel jobs, instead of submitting (do not use on compute canada login node)"
   echo " -A <account> : account to use for allocation (default: $CC_COMPUTE_ALLOC)"
   echo ""
   usage_job_templates
   usage_job_depends

 }


function getContainer {

app_name=$1
bids_app_dir=$SINGULARITY_DIR/bids-apps

linenum=1
foundmatch=0
for app in `awk -F '\t' '{print $1}' $app_list`
do

 if [ "$app" = "$app_name" ]
 then
	foundmatch=1
	break
 fi
 linenum=$((linenum+1))	
done


if [ "$foundmatch" = 0 ]
then
	echo "${app_name} does not exist in bids-app list!"  >&2 
	exit 1
fi

url=`cat $app_list | head -n $linenum | tail -n 1 |  awk -F '\t' '{print $2}'`
base_opts=`cat $app_list | head -n $linenum | tail -n 1 |  awk -F '\t' '{print $3}'`
app_job_template=`cat $app_list | head -n $linenum | tail -n 1 |  awk -F '\t' '{print $4}'`

hub=${url%%:*}
if [ "$hub" = "docker" ]
then

  #use shub-cache to download if not exist, and return path to image
 singularity_image=`shub-cache $url`
 if [ ! "$?" = 0  -o ! -n "$singularity_image" -o ! -e "$singularity_image"  ]
 then
   echo "Failed to pull $url, exiting" >&2
   exit 1
 fi



fi

if [ "$hub" = "shub" ]
then
 #use shub-cache to download if not exist, and return path to image
 singularity_image=`shub-cache $url`
 if [ ! "$?" = 0  -o ! -n "$singularity_image" -o ! -e "$singularity_image"  ]
 then
   echo "Failed to pull $url, exiting" >&2
   exit 1
 fi
 

fi

if [ "$hub" = "file" ]
then
 singularity_image=${url##file://}

 if [  ! -e "$singularity_image"  ]
 then
   echo "File $singularity_image does not exist,  exiting" >&2
   exit 1
 fi
 

fi


if [ "$override_template" = "0" -a -n "$app_job_template" ]
then
	job_template=$app_job_template 
	echo "Using job-template specified in bids-apps.tsv, $app_job_template" >&2
fi




}


app_list=$NEUROGLIA_DIR/bids-apps.tsv

if [ ! -n "$SINGULARITY_OPTS" ]
then
	echo "SINGULARITY_OPTS not defined! exiting..." >&2
	exit 1
fi

cmd=$0
cmdname=${cmd##*/}


#add option for this
#	   #get app usage
#	   echo getting app usage
#	   singularity run -e $singularity_image


job_template=Regular
depends=""
singlesubj=""
testmode=0
subjlist=""
override_template=0
runlocal=0
localscratch="/localscratch" # default for graham
cc_account=$CC_COMPUTE_ALLOC

while getopts "Jj:Bd:s:tS:L:A:" options; do
 case $options in
    B ) print_app_list_long
	    exit 1;;
    J ) print_job_templates
	exit 1;;
    j ) echo "	Using job template: $OPTARG" >&2
	override_template=1
	job_template=$OPTARG;;
    d ) echo "	Using dependencies: $OPTARG" >&2
	depends=$OPTARG;;
    S ) echo "	Using subject list: $OPTARG" >&2
	subjlist=$OPTARG;;
    s ) echo "	Using single subject: $OPTARG" >&2
	singlesubj=$OPTARG;;
    t ) echo "	Using test-mode (no submit jobs)" >&2
	testmode=1;;
    L ) echo "  Using local mode (run jobs locally instead of submitting)" >&2
    ncores=$OPTARG
    host=`hostname`
    if [ ! "${host:0:3}" = "gra" ]  #if not running on graham, then:
    then
    localscratch=/tmp
    SINGULARITY_OPTS="$SINGULARITY_OPTS -e -B $localscratch:/scratch"
    fi
    runlocal=1;;
    A ) echo "	Using allocation account: $OPTARG" >&2
	    cc_account=$OPTARG;;
    * ) usage
	exit 1;;
 esac
done

shift $((OPTIND-1))

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


app_name=$1

shift 1

getContainer $app_name



if [ "$#" -lt 3 ]
then
	usage
	
	echo "$app_name, built-in default options: $base_opts"
	echo "Usage for bids app:" >&2
	if [ -e $singularity_image ]
	then
		singularity run -e $singularity_image
	fi

	
	exit 1
fi

bids_dir=`realpath $1`
out_dir=$2
level=$3

#check validity of participant 
level_part=${level:0:5}
if [   "$level_part" != "group" -a  "$level_part" != "parti" ]
then
	echo "Required argument 3 to bids apps must be analysis level (participant/group)" >&2
	exit 1
fi

mkdir -p $out_dir
out_dir=`realpath $out_dir`

shift 3
options=$@


#check if participant_label is used as an option 
#and quit if it is found

if [[ $options = *--participant_label* ]]
then
  echo "ERROR: --participant_label should not be used with bidsBatch, use: bidsBatch -s <participant_label> instead" >&2
 exit 1
fi  

job_dir=$out_dir/code/jobs
mkdir -p $job_dir


list=$bids_dir/participants.tsv

if [ ! -e $list ]
then
 list=$job_dir/in_participants.tsv
 echo "$list does not exist, creating temporary one in $list" >&2 
 echo "participant_id"> $list
 pushd $bids_dir >&2 
 ls -d sub-* >> $list
 popd  >&2 

fi

#check if header line is there
header=`head -n 1 $list`
if [ ! "${header:0:14}" = "participant_id" ]
then
	echo "invalid $list, no header row.."
       	list=$job_dir/in_participants.tsv
	echo "$list invalid (no header row), creating temporary one in $list" >&2 
	 echo "participant_id"> $list
	 pushd $bids_dir >&2
	 ls -d sub-* >> $list
	 popd >&2

fi

N=`tail -n +2 $list | wc -l`

if [ ! -n "$singlesubj" ]
then
	if [ ! -n "$subjlist" ]
	then
		indices="1-$N"  # all lines in subjlist, 1 to N
	else
		#pick out indices from subjlist
		indices=""
		for sub in `cat $subjlist`
		do
			N_matched=`grep -n sub-$sub $list | wc -l` 
			ind=`grep -n sub-$sub $list`
			if [ "$N_matched" = 1 ]
			then
				ind=${ind%%:*}
				if [ -n "$indices" ]
				then
					indices="$indices,$((ind-1))"
				else
					indices="$((ind-1))"
				fi
			fi
		done
	fi
else
	N_matched=`grep -n sub-$singlesubj $list | wc -l` 
	if [ "$N_matched" == 1 ]
	then
		ind_subj=`grep -n sub-$singlesubj $list`
		indices=${ind_subj%%:*}
		indices=$((indices-1))
	fi

	if [ "$N_matched" == 0 ]
	then
		echo "single_subj: $singlesubj does not match any subjids in $list!" >&2
		exit 1
	fi

	if [ "$N_matched" -gt 1 ]
	then
		echo "single_subj: $singlesubj matches multiple subjids in $subjlist!" >&2
		exit 1
	fi
fi


job=$job_dir/$app_name.$level.$RANDOM.job

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


#forces number of threads - based on job template
ncpus=`grep cpus-per-task $job`
ncpus=${ncpus##*=}
FORCE_THREADS="export SINGULARITYENV_OMP_NUM_THREADS=$ncpus; export SINGULARITYENV_MKL_NUM_THREADS=$ncpus;"

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

if [ "$level_part" == "parti" ]
then
	#submit with array, using subj and subjlist, indexed by line number (1-N)
	echo "#SBATCH --output=$job_dir/${app_name}.$level.%A_%a.out" >> $job
	echo "#SBATCH --array=$indices" >> $job
	echo "$FORCE_THREADS" >> $job
	echo "subj=\`tail -n +2 `realpath $list` | head -n \$SLURM_ARRAY_TASK_ID  | awk '{print \$1}' | tail -n 1\`" >> $job
	#strip off sub- from sub-LABEL
	echo "subj=\${subj##sub-}" >> $job
	echo "echo SINGULARITY_IMAGE: $singularity_image" >> $job
	echo "echo CMD: $bids_dir $out_dir $level --participant_label \$subj $base_opts $options" >>$job
	echo "echo START_TIME: \`date\`" >>$job
	echo "cd `pwd`" >> $job
	echo singularity run -B $localscratch:/localscratch $SINGULARITY_OPTS $singularity_image $bids_dir $out_dir $level  --participant_label \$subj $base_opts $options >> $job
	echo "RETURNVAL=\$?" >> $job
	echo "echo RETURNVAL=\$RETURNVAL" >>$job
	echo "echo END_TIME: \`date\`" >>$job
	echo "exit \$RETURNVAL" >> $job

	echo -n "		Creating participant-level job ... " >&2
else
	#submit without arrays, for group level
	echo "#SBATCH --output=$job_dir/${app_name}.$level.%A.out" >> $job
	echo "$FORCE_THREADS" >> $job
	echo "echo SINGULARITY_IMAGE: $singularity_image" >> $job
	echo "echo CMD:  $bids_dir $out_dir $level $base_opts $options" >>$job
	echo "echo START_TIME: \`date\`" >>$job
	echo "cd `pwd`" >> $job
	echo singularity run -B $localscratch:/localscratch $SINGULARITY_OPTS $singularity_image $bids_dir $out_dir $level $base_opts $options >> $job
	echo "RETURNVAL=\$?" >> $job
	echo "echo RETURNVAL=\$RETURNVAL" >>$job
	echo "echo END_TIME: \`date\`" >>$job
	echo "exit \$RETURNVAL" >> $job

	
	echo -n "		Creating group-level job  ... " >&2
fi


if [ -n "$depends" ]
then
	dependsopt="--dependency=$depends"
fi

if [ "$testmode" == 0 ] && [ "$runlocal" == 0 ]
then
echo "Submitting job" >&2 
message=`sbatch $dependsopt $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/$app_name.$level.$jobid.job

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

else #either running in testmode or locally
    if [ "$testmode" == 1 ]
    then
	    jobid=$RANDOM
    	echo "fake-jobid=$jobid  (test-mode: no jobs submitted)" >&2
    else
        if [ "$runlocal" == 1 ]
        then

        popts="--no-notice --verbose --jobs $ncores --line-buffer "

        #if no special cases, then make a run script, then run that script
        if [ ! -n "$subjlist" ]
        then
   
        datestring=`date +%Y-%m-%d_%Hh%Mm`
        joblist=$job_dir/joblist.local_$datestring.txt
        if [ "$level_part" == "parti" ]
        then
            if [ ! -n "$singlesubj" ]
            then
                 seq 1 $N | parallel $popts echo "export SLURM_ARRAY_TASK_ID={}\; bash $job 2\>\&1 \| tee $job_dir/${app_name}.$level.local_$datestring.{}.out" > $joblist
                 else
                 echo $indices | parallel $popts echo "export SLURM_ARRAY_TASK_ID={}\; bash $job 2\>\&1 \| tee  $job_dir/${app_name}.$level.local_$datestring.{}.out" > $joblist
            fi
         else
            #group level:
           echo "bash $job 2\>\&1 \| tee  $job_dir/${app_name}.$level.local_$datestring.out" > $joblist
        fi

        else
            echo "subjlist mode not implemented yet with -L option.." >&2
            exit 1
        fi
       
        echo "Running jobs locally..." >&2
        parallel $popts :::: $joblist


    fi

    fi

    
fi

echo "Job and job output stored in: $job_dir" >&2
echo $jobid

exit 0
