#!/bin/bash

################################################################################
## Launcher script for SUNSET
################################################################################
##
## It reads a SUNSET recipe and splits it into single ('atomic') verifications.
## Then each atomic verification is launched as an independent job with SLURM.
## Each job can be further parallelized using the --cpus parameter and the 
## 'ncores' field in the recipe. 
## If requested in the recipe, the Scorecards for the results are also created
## and saved once the verifications are finished. 
##
################################################################################

# Usage statement

function usage
{
  echo "Usage: $0 <path_to_recipe> <path_to_script> --wallclock=<hh:mm:ss> --custom_directives=<custom_slurm_directives> --disable_unique_ID"
  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'."
  echo " "
  echo " --disable_unique_ID: Do not add a unique ID to the output folder name."
}

if [[ ( $@ == "--help") || $@ == "-h" ]]; then
  usage
  exit 0
fi

# 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
      ;;
    --disable_unique_ID)
      export disable_unique_ID="--disable_unique_ID"
      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
# module load conda/22.11.1-2
# conda activate condaCerise
# export LD_LIBRARY_PATH=/perm/cyce/conda/envs/condaCerise/lib

Rscript split.R ${recipe} $disable_unique_ID --tmpfile $tmpfile

# Run with Autosubmit or directly with Slurm's sbatch?
run_method=$( head -1 $tmpfile | tail -1 )
# If run method is 'sbatch', launch jobs with dependencies
if [[ $run_method == "sbatch" ]]; then
  # Retrieve working directory
  codedir=$( head -2 $tmpfile | tail -1 )
  # Retrieve output directory
  outdir=$( head -3 $tmpfile | tail -1 )
  # Multimodel TRUE/FALSE
  multimodel=$( head -4 $tmpfile | tail -1)
  # Scorecards TRUE/FALSE
  scorecards=$( head -5 $tmpfile | tail -1)
  
  # Create directory for slurm output
  logdir=${outdir}/logs/slurm/
  mkdir -p $logdir
  echo "Slurm job logs will be stored in $logdir"

  # Get corrected recipe
  recipe=${outdir}/logs/recipes/$(basename $recipe)

  # Is machine MN5?
  if [[ $BSC_MACHINE == "mn5" ]]; then
    platform_params="-A bsc32 -q gp_bsces"
  else
    platform_params=""
  fi
  
  # Launch one job per atomic recipe
  cd $codedir
  job_number=0
  # Create empty array to store all the job IDs for the verification jobs
  verification_job_list=()
  echo "Submitting verification jobs..."
  # Loop over atomic recipes
  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
    # Send batch job and capture job ID
    job_ID=$(sbatch --parsable --job-name="SUNSET_verification" $platform_params --nodes=1 --output=$outfile --error=$errfile --time=$wallclock --cpus-per-task=$cpus $custom_directives conf/slurm_templates/run_parallel_workflow.sh ${script} ${atomic_recipe})
    # Add job ID to array
    verification_job_list+=($job_ID)
    echo "Submitted batch job $job_ID"
  done

  multimodel_job_list=()
  job_number=0
  if [[ $multimodel == "TRUE" ]]; then
    for atomic_recipe in ${outdir}/logs/recipes/multimodel/atomic_recipe_*.yml; do
      job_number=$(($job_number + 1))
      job_name=$(basename $outdir)_$(printf %02d $job_number)
      outfile=${logdir}/run-multimodel-${job_name}.out
      errfile=${logdir}/run-multimodel-${job_name}.err
      # Send batch job and capture job ID
      job_ID=$(sbatch --parsable --dependency=afterok:$(IFS=,; echo "${verification_job_list[*]}") --kill-on-invalid-dep=yes --job-name="SUNSET_multimodel" $platform_params --nodes=1 --output=$outfile --error=$errfile --time=$wallclock --cpus-per-task=$cpus $custom_directives conf/slurm_templates/run_parallel_workflow.sh ${script} ${atomic_recipe})
      # Add job ID to array
      multimodel_job_list+=($job_ID)
      echo "Submitted batch job $job_ID"
    done
  fi
  
  # Submit scorecards job with dependency on verification jobs, passed as a 
  # comma-separated string. The scorecards  job will not run until all the 
  # verification jobs have finished successfully.
  # If any of the jobs fail, it will be canceled.
  if [[ $scorecards == "TRUE" ]]; then
    echo "Submitting scorecards jobs..."
    outfile=${logdir}/run-scorecards.out
    errfile=${logdir}/run-scorecards.err
    sbatch --dependency=afterok:$(IFS=,; echo "${verification_job_list[*]} ${multimodel_job_list[*]}") $platform_params --output=$outfile --error=$errfile --time=01:00:00 conf/slurm_templates/run_scorecards.sh ${recipe} ${outdir}
  fi
fi

# Clean temporary file
rm $tmpfile
