#!/bin/bash

while (("$#")); do
  case "$1" in
  -h | --help)
    echo "Usage: $(basename "$0") [OPTION]..."
    echo ""
    echo "Submit a slurm batch job for the given task and other options."
    echo ""
    echo "  -h, --help                     Print this message"
    echo "  -s, --single-objective         Single-objective search (task performance only)"
    echo "  -t, --task TASK                The task/dataset name for which to run the"
    echo "                                 experiment (required)"
    echo "  -r, --reservation RESERVATION  The reservation name, if applicable"
    echo "  -n, --reservation-nodes RESERVATION_NODES"
    echo "                                 The reserved number of nodes, if applicable"
    exit 0
    ;;
  -s | --single-objective)
    if [ -n "$SINGLE_OBJECTIVE" ]; then
      echo "Error: flag $1 specified multiple times" >&2
      exit 2
    fi
    SINGLE_OBJECTIVE="true"
    shift
    ;;
  -t | --task)
    if [ -n "$TASK" ]; then
      echo "Error: flag $1 specified multiple times" >&2
      exit 2
    elif [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
      TASK=$2
      shift 2
    else
      echo "Error: argument for $1 is missing" >&2
      exit 2
    fi
    ;;
  -r | --reservation)
    if [ -n "$RESERVATION" ]; then
      echo "Error: flag $1 specified multiple times" >&2
      exit 2
    elif [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
      RESERVATION=$2
      shift 2
    else
      echo "Error: argument for $1 is missing" >&2
      exit 2
    fi
    ;;
  -n | --reservation-nodes)
    if [ -n "$RESERVATION_NODES" ]; then
      echo "Error: flag $1 specified multiple times" >&2
      exit 2
    elif [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
      RESERVATION_NODES=$2
      shift 2
    else
      echo "Error: argument for $1 is missing" >&2
      exit 2
    fi
    ;;
  *)
    echo "$(basename "$0"): invalid option -- $1" >&2
    echo "Try '$(basename "$0") --help' for more information." >&2
    exit 2
    ;;
  esac
done

if [ -z "$TASK" ]; then
  echo "Error: missing required argument: -t, --task TASK" >&2
  echo "Try '$(basename "$0") --help' for more information." >&2
  exit 2
fi

if [ -n "$RESERVATION" ] && [ -z "$RESERVATION_NODES" ]; then
  echo "Error: if a reservation, -n, --reservation-nodes RESERVATION_NODES must" >&2
  echo "also be specified." >&2
  echo "Try '$(basename "$0") --help' for more information." >&2
  exit 2
fi

if ! command -v sbatch &> /dev/null; then
  echo "You do not appear to be on a slurm-based job submission system!" >&2
  echo "Details: sbatch is not a command on \$PATH." >&2
  echo "Exiting without submitting a job :(" >&2
  exit 1
fi


trap rm_tmp INT

function rm_tmp() {
  rm -rf "$batch_script_dir"

  exit 130
}


batch_script_dir=$(mktemp --directory)
batch_script="$batch_script_dir/xnas_batch_script_$TASK.sh"


cat > "$batch_script" <<- EOM
#!/bin/bash
#SBATCH --output=logs_batch/slurm-xnas-$TASK-%A_%a.out
#SBATCH --job-name=xnas-$TASK
EOM


if [ -n "$RESERVATION" ]; then
  cat >> "$batch_script" <<- EOM
# start: reservation-specific directives
#SBATCH --array=0-$(( RESERVATION_NODES-1 ))
#SBATCH --time=0
#SBATCH --qos=exempt
#SBATCH --reservation=$RESERVATION
# end: reservation-specific directives
EOM
else
  cat >> "$batch_script" <<- EOM
#SBATCH --array=0-15
#SBATCH --time=1-00:00:00
EOM
fi

cat >> "$batch_script" <<- EOM
#SBATCH --partition=pbatch
EOM

cat >> "$batch_script" <<- EOM
echo "Started at \$(date)"
# symlink to "latest" log file
ln -fs \
"slurm-xnas-$TASK-\${SLURM_ARRAY_JOB_ID}_\${SLURM_ARRAY_TASK_ID}.out" \
"logs_batch/slurm-xnas-$TASK-latest_\${SLURM_ARRAY_TASK_ID}.out"
env | grep SLURM
srun ./xnas-ray-job.sh $TASK $SINGLE_OBJECTIVE
status=\$?
echo "Finished at \$(date) with status \$status"
exit \$status
EOM

echo "Submitting the following job:"
echo "============================================================================="

cat "$batch_script"
echo "============================================================================="

sbatch "$batch_script"


rm_tmp
