#!/bin/bash
set -e
set -o pipefail


RED='\033[0;31m'
GREEN='\033[0;32m'
NC='\033[0m' # No Color


g() {
  echo -e -n '\033[0;32m'
};
n() {
  echo -e -n '\033[0m'  # No Color
}

ECHO_DIR=$(pwd)
BRC_DIR=${ECHO_DIR}/brc
if [[ $ECHO_DIR != */echo ]]; then
  echo -e "${RED} Please run this script from the 'echo' directory ${NC}";
  return 0 2> /dev/null || exit 0
fi;
#shellcheck disable=SC2162
# shellcheck disable=SC1073

gen_jobs() {
    g
    echo -e "Pick an option:"
    n
    options=("Specify a script to generate jobs.json file (i.e: scripts/single, scripts/all) OR"
    "Specify an experiment folder (must have run create_experiment_params first) OR"
    "Help")

    select opt in "${options[@]}"; do
        case "$REPLY" in
            1 ) g; echo -n "Found these scripts: "; n; echo $(ls scripts | sed 's/^/scripts\//;s/$/\ , /') ;g;
                   echo  -n "Path of script: ";n; read -e script;
                jobs=$(bash $script make | tail -1); break;;
            2 ) g; echo -n "Path of experiment folder (i.e experiments/gradient_passing/QPSK_neural_and_neural): ";n; read -e ef;
                echo "RUNNING: python ${ECHO_DIR}/experiments/make_jobs.py --experiment_folder=${ef}..."
                python "${ECHO_DIR}/experiments/make_jobs.py --experiment_folder=${ef}" ; jobs=$ef; break;;
            3 ) echo "try running python experiments/create_experiment_params.py to make an experiment";
                exit 0;;
            *)  echo "Invalid option. Try another one.";continue;;
        esac
    done
}
mkdir -p _tmp_
g
echo -e "Are you on the BRC cluster?"; n;
select yn in "Yes" "No"; do
    case $yn in
        Yes ) ONBRC=true ; echo "You are on the BRC.";
              module load python/3.6 gcc; python=python3 ;
              #source activate runecho;
              break;;
        No )  ONBRC=false ; echo "You are not on the BRC.";  break;;
        *) echo "Invalid option. Try another one.";continue;;
    esac
done

g
echo -e "Do you already have a jobs.json file?";n;
select yn in "Yes" "No"; do
    case $yn in
        Yes ) break;;
        No ) gen_jobs; break;;
        *) echo "Invalid option. Try another one.";continue;;
    esac
done

# shellcheck disable=SC2162
if [[ -z ${jobs+x} ]]; then
  g; echo -n "Specify jobs.json file (or directory containing jobs.json) to process: " ; n ; read -e jobs;
else
  g; echo -n "FOUND:";n;  echo -n "${jobs}"; g; echo "... is this the correct jobs json file?"; n;
  select yn in "Yes" "No"; do
      case $yn in
          Yes ) break;;
          No ) g; echo -n "Specify jobs.json file (or directory containing jobs.json) to process: " ; n ; read -e jobs; break;;
          *) echo "Invalid option. Try another one.";continue;;
      esac
  done
fi;

if [[ $jobs != *.json ]] ; then
  jobsjson=${jobs}jobs.json
else
  jobsjson=${jobs}
fi

# shellcheck disable=SC2162
# shellcheck disable=SC2207

if [[ "$ONBRC" = true ]] ; then
  echo "!!! ON BRC !!!"
  mkdir -p "${BRC_DIR}/out"
  while true; do
      g; echo -n "How many nodes do you want to use? [1-8]: ";n; read -e nodes;
      if ! [[ "$nodes" =~ ^[0-9]+$ ]] ; then
          echo "Not an integer";
          continue;
      fi;
      if [[ $nodes -gt 8 || $nodes -lt 1 ]] ; then
          echo "Invalid number of nodes: $nodes";
          continue;
      fi;
      if [[ $nodes == 0 ]] ; then
          echo "Exiting, Goodbye!";
          exit 0;
      fi;
      break;
  done;
#  source deactivate; #dont need jsonnet anymore
  tasks=($(python ${BRC_DIR}/make_tasks.py ${jobsjson} $nodes | tail -$((3 * nodes))))
#  echo "python ${BRC_DIR}/make_tasks.py ${jobsjson} $nodes | tail -$((2 * nodes))"
  for i in $(eval echo "{0..$((nodes-1))}"); do
      taskfile=${tasks[$((2 * i))]}
      n=${tasks[$(((2 * i)+1))]}
      echo "$taskfile $n"
      mkdir -p ${ECHO_DIR}/brc/out/work
      sbatch ${ECHO_DIR}/brc/SLURM_tasks.sh ${taskfile} ${n}
  done
  # sbatch "${ECHO_DIR}/brc/SLURM_tasks.sh ${taskfile} ${numtasks}"
else
  g;
  echo "Would you like to run [all/one] job(s) from $jobsjson?";
  n;
  select howmany in "all" "one"; do
      case $howmany in
          all ) python ${ECHO_DIR}/run_experiment.py --jobs_file=${jobsjson}; break;;
          one ) g; jid=0; echo -n "Specify the job_id of the job you want to run: ";n; read -e jid;
              python ${ECHO_DIR}/run_experiment.py --jobs_file=${jobsjson} --job_id=${jid} --plot; break;;
          *) echo "Invalid option. Try another one.";continue;;
      esac
  done
fi;

function finish {
  n;
  set +e
}
trap finish EXIT