#!/bin/bash 
source /etc/profile.d/modules.sh

#set -eu
#set -v
set -e

function usage {
  echo "Usage: $0 -e=<EXPID> -r=<recipe_template> -h=<horizon> [-cd=<code_dir>] [-sd=YYYYMMDD] [--rr=no]"
  echo " "
  echo "  -e, --expid:       Autosubmit EXPID"
  echo "  [MANDATORY]"
  echo "  -r, --recipe:      Path to YAML recipe template"
  echo "  [MANDATORY]"
  echo "  -h, --horizon:     Forecast horizon: 'subseasonal', 'seasonal', or 'decadal'"
  echo "  [MANDATORY]"
  echo "  -sd, --start_date: Start date (YYYYMMDD) of the forecast."
  echo "                     If not specified, it will be set as the"
  echo "                     first day of the current month (seasonal) or the"
  echo "                     most recent Thursday (subseasonal)."
  echo "  -cd, --code_dir:   Directory where the SUNSET code is stored."
  echo "                     The default value is the current working directory."
  echo "  -rr, --rerun:      Options: 'no' (default), 'rerun_failed', 'skip_failed'"
  echo "                     - 'no': The Autosubmit experiment is newly generated"
  echo "                     for the specified start date, and the workflow"
  echo "                     runs from the beginning as normal."
  echo "                     - 'rerun_failed': The FAILED jobs in the Autosubmit"
  echo "                     experiment are run again, for the same start date"
  echo "                     previously created."
  echo "                     - 'skip_failed': The FAILED jobs in the Autosubmit"
  echo "                     experiment are skipped so that any pending workflow"
  echo "                     dependencies can run. Dependencies might fail if"
  echo "                     they require outputs from the FAILED jobs."
  echo " "
}


function parse_args {
  for i in "$@"; do
    case $i in
      -e=*|--expid=*)
        export expid=`echo $1 | sed -e 's/^[^=]*=//g'`
        shift
        ;;
      -r=*|--recipe=*)
	export recipe_template=`echo $1 | sed -e 's/^[^=]*=//g'`
	shift
	;;
      -sd=*|--start_date=*)
        export start_date=`echo $1 | sed -e 's/^[^=]*=//g'`
        shift
        ;;
      -cd=*|--code_dir=*)
        export code_dir=`echo $1 | sed -e 's/^[^=]*=//g'`
        shift
        ;;
      -h=*|--horizon=*)
        horizon="${i#*=}"
        shift # past argument=value
        ;;
      -r=*|--rerun=*)
        export rerun=`echo $1 | sed -e 's/^[^=]*=//g'`
        shift
        ;;
      -h|--help)
        usage
        shift # past argument with no value
        ;;
      *)
      ;;
    esac
  done
  if [ -z "$start_date" ]; then
    start_date=$(date +%Y%m%d)
  fi
  if [ -z "$code_dir" ]; then
    code_dir=./
  fi
  if [ -z "$rerun" ]; then
    rerun='no'
  fi
}

function get_startdate {
  for dl_date in ${start_date[@]}; do 
    if [[ $horizon == "subseasonal" ]];then
      week_day=$(date --date="$dl_date" +%a)
      # Check date and stream
      if [[ ${week_day} == Thu ]]; then
        sdate=($dl_date)
      else
        sdate=$(date --date "last Thursday" +%Y%m%d)
      fi
    else
      sdate=${dl_date:0:6}01
    fi
  done
}

function launch_auto {

  rerun=${rerun[@]}
  username=$(whoami)
  if [ $rerun == "no" ]; then
    code_dir=${code_dir[@]}
    echo "*** Launching operational: ***"
    echo "[INFO] CURRENT TIME: $(date)"
    echo "[INFO] STARTDATE: ${sdate[@]}"
    echo "[INFO] Using EXPID: ${expid}"
    echo "[INFO] Running SUNSET launcher from ${code_dir}" 

    cd ${code_dir}
    source MODULES
    launcher_args=$(Rscript recipe_maker.R ${recipe_template} --sd=${sdate[@]} --horizon=${horizon} --expid=${expid})
    bash ${code_dir}/launch_SUNSET.sh ${launcher_args} --disable_unique_ID

    ## TODO: Set up ssh to AS machine; or do it directly as part of launch_SUNSET.sh?
    ssh ${username}@bscesautosubmit04.bsc.es "module load autosubmit/4.1.14-foss-2023b-Python-3.11.5;
      autosubmit create ${expid} --noplot;
      autosubmit refresh ${expid};
      # create must be done each time 
      autosubmit monitor ${expid};
      # Save for logging purposes (specify -fs to filter
      # jobs when ensemble, -o {pdf,png,ps,svg} for different plot format files)
      nohup autosubmit run -v ${expid} &"
  elif [ $rerun == "rerun_failed" ]; then
    echo "*** Re-running failed jobs ***"
    echo "[INFO] CURRENT TIME: $(date)"
    echo "[INFO] STARTDATE: ${sdate[@]}"
    echo "[INFO] EXPID: ${expid}"
    ssh ${username}@bscesautosubmit04.bsc.es "module load autosubmit/4.1.14-foss-2023b-Python-3.11.5;
      autosubmit setstatus ${expid} -np -fs FAILED -t WAITING -s;
      nohup autosubmit run -v ${expid} &"
  elif [ $rerun == "skip_failed" ]; then
    echo "*** Skipping failed jobs ***"
    echo "[WARNING] Previously failed jobs will NOT be completed"
    echo "[INFO] CURRENT TIME: $(date)"
    echo "[INFO] STARTDATE: ${sdate[@]}"
    echo "[INFO] EXPID: ${expid}"
    ssh ${username}@bscesautosubmit04.bsc.es "module load autosubmit/4.1.14-foss-2023b-Python-3.11.5;
      autosubmit setstatus ${expid} -np -fs FAILED -t COMPLETED -s;
      nohup autosubmit run -v ${expid} &"
  fi 

}

#########
# MAIN 
#########

parse_args $@
# Retrieve startdate
get_startdate
# Launch Workflow 
launch_auto
