#!/bin/bash

# Usage statement
if [[ ( $@ == "--help") || $@ == "-h" ]]; then
  echo "Usage: $0 <path_to_recipe> <path_to_script> --wallclock=<hh:mm:ss> --custom_directives=<custom_slurm_directives>"
  exit 0
fi

function usage
{
  echo "Usage: $0 <path_to_recipe> <path_to_script> --wallclock=<hh:mm:ss> --custom_directives=<custom_slurm_directives>"
  echo " "
  echo "  <path_to_recipe>: Path to the SUNSET recipe."
  echo " "
  echo "  <path_to_script>: Path to the user-defined script."
  echo " "
  echo " --wallclock=<hh:mm:ss>: Maximum execution time for the jobs."
  echo "                         Default is 02:00:00."
  echo " "
  echo " --cpus=<number_of_cpus>: CPUs (cores) to be requested for each job."
  echo "                          Corresponds to the --cpus-per-task parameter"
  echo "                          in sbatch. The default number is 1."
  echo " "
  echo " --custom_directives=<custom_directives>: Custom directives for sbatch."
  echo "                                          E.g. '--constraint=medmem'."
}


# Assign arguments
recipe=$1
script=$2

for i in "$@"; do
  case $i in
    --wallclock*)
      export wallclock=`echo $1 | sed -e 's/^[^=]*=//g'`
      shift
      ;;
    --cpus*)
      export cpus=`echo $1 | sed -e 's/^[^=]*=//g'`
      shift
      ;;
    --custom_directives*)
      export custom_directives=`echo $1 | sed -e 's/^[^=]*=//g'`
      shift
      ;;
    -h|--help)
      usage
      exit 0
      shift
      ;;
    *)
      shift
      ;;
  esac
done

# Check recipe
if [ ! -f "$recipe" ]; then
  echo "Could not find the recipe file: $recipe"
  usage
  exit 1
fi
# Check script
if [ ! -f "$script" ]; then
  echo "Could not find the script file: $script"
  usage
  exit 1
fi
# Assign default wallclock time if not specified
if [ -z "$wallclock" ]; then
  wallclock="02:00:00"
fi
# Assign default number of cores if not specified
if [ -z "$cpus" ]; then
  cpus=1
fi
# Assign empty custom directives if not specified
if [ -z "$custom_directives" ]; then
  custom_directives=""
fi

# Define tmp file to store necessary information
tmpfile=$(mktemp ${TMPDIR-/tmp}/SUNSET.XXXXXX)

# Create outdir and split recipes
source MODULES
Rscript split.R ${recipe} --tmpfile $tmpfile
# Retrieve working directory
codedir=$( head -1 $tmpfile | tail -1 )
# Retrieve output directory
outdir=$( head -2 $tmpfile | tail -1 )

rm $tmpfile

# Create directory for slurm output
logdir=${codedir}/out-logs/slurm_logs/
mkdir -p $logdir
echo "Slurm job logs will be stored in $logdir"

# Launch one job per atomic recipe
cd $codedir
job_number=0
for atomic_recipe in ${outdir}/logs/recipes/atomic_recipe_??.yml; do
  job_number=$(($job_number + 1))
  job_name=$(basename $outdir)_$(printf %02d $job_number)
  outfile=${logdir}/run-${job_name}.out
  errfile=${logdir}/run-${job_name}.err
  sbatch --output=$outfile --error=$errfile --time=$wallclock --cpus-per-task=$cpus $custom_directives conf/slurm_templates/run_parallel_workflow.sh ${script} ${atomic_recipe}
done

## TODO: Retrieve jobids and list them as dependencies for scorecards job?
