#!/bin/bash
# Submit round 5 runs: per run, the dynamic melt_f calibration (40 chunks of
# 100 glaciers, keeping the calibration run), then the merge.
#
#   ./submit.sh <run> [<run> ...]          # e.g. default pf1.5 default_hugp
#   ./submit.sh --test <run>               # the 4 TEST_IDS glaciers, from the
#                                          # exp's test_run/ gdirs
#   AFTER=<jobid> ./submit.sh <run>        # wait for that job (the gdir build)
#   DYN_CHUNKS=3,17 ./submit.sh <run>      # redo those chunks, then the merge
#   ONLY=merge ./submit.sh <run>           # just the merge
#
# <run> is a folder of ../lmr_gdirs_v7_202610/exps, optionally with a suffix:
#   <exp>_hugp   Hugonnet target + HUG_SIGMA (less mass loss)
#   <exp>_hugm   Hugonnet target - HUG_SIGMA (more mass loss)
#   ..._ys<year> calibration runs start in <year> instead of DYN_YS
#   ..._scan     scan + bisection calibration (scan_calib.py), last suffix
#                (e.g. default_ens148_ys1500, default_hugp_ys1500)
#
#   output/<run>/chunks/                      per-chunk files
#   output/<run>/hist_dyn_calib_reg11.nc      the calibration run, DYN_YS-2025
#   output/<run>/dyn_calib_reg11.csv          per-glacier calibration table
#   output/<run>/lmr_mira_hist_reg11.nc       only if RUN_YS is set
set -e

HERE=$(cd "$(dirname "$0")" && pwd)
source "$HERE/common.sh"

TEST=""
OUT_DIR=$R5_DIR/output
LOGDIR=$R5_DIR/_logs
DYN_CHUNKS_ALL=0-$(( RUN_N_CHUNKS - 1 ))
if [ "$1" = "--test" ]; then
    shift
    TEST=test
    OUT_DIR=$R5_DIR/test_run
    LOGDIR=$R5_DIR/_logs_test
    # the chunks of 100 holding the TEST_IDS (11.00897 -> 8, ...)
    DYN_CHUNKS_ALL=$(for i in $TEST_IDS; do echo $(( 10#${i#RGI60-11.} / RUN_CHUNK_SIZE )); done | paste -sd,)
    DYN_RUN_TIME=02:00:00
fi

if [ -z "$1" ]; then
    echo "Usage: $0 [--test] <run> [<run> ...]" >&2
    exit 1
fi

# <run> -> the gdir experiment and the Hugonnet shift
parse_run () {
    # optional _ys<year>: start year of the calibration runs (default DYN_YS)
    local r=$1
    # optional _scan: scan + bisection calibration (scan_calib.py)
    CALIB_RUN=
    if [[ $r =~ ^(.+)_scan$ ]]; then r=${BASH_REMATCH[1]}; CALIB_RUN=scan; fi
    YS_RUN=
    if [[ $r =~ ^(.+)_ys([0-9]+)$ ]]; then r=${BASH_REMATCH[1]}; YS_RUN=${BASH_REMATCH[2]}; fi
    set -- "$r"
    case $1 in
        *_hugp) EXP=${1%_hugp}; SHIFT=$HUG_SIGMA ;;
        *_hugm) EXP=${1%_hugm}; SHIFT=-$HUG_SIGMA ;;
        *)      EXP=$1;         SHIFT= ;;
    esac
}

for RUN in "$@"; do
    parse_run "$RUN"
    EXP_DIR=$V7_EXPS_DIR/$EXP
    # the experiment's border (config.sh, in a subshell: it re-sources common.sh)
    BDIR=$(source "$EXP_DIR/config.sh"; printf "b_%03d" "$BORDER")
    GDIRS=$EXP_DIR/RGI62/$BDIR/L3
    [ -n "$TEST" ] && GDIRS=$EXP_DIR/test_run/RGI62/$BDIR/L3
    if [ ! -f "$EXP_DIR/config.sh" ]; then
        echo "ERROR: no v7 experiment '$EXP' (run '$RUN')" >&2; exit 1
    fi
    if [ -z "$AFTER" ] && [ "$ONLY" != "merge" ] && [ ! -d "$GDIRS" ]; then
        echo "ERROR: no gdirs for '$EXP': $GDIRS (build first, or AFTER=<jobid>)" >&2
        exit 1
    fi
done

echo "------------------------------------------------------------------"
[ -n "$TEST" ] && echo "  *** TEST RUN: $TEST_IDS"
echo "  runs          $*"
echo "  output        $OUT_DIR/<run>/"
echo "  dyn           from $DYN_YS, Hugonnet err x $DYN_ERR_FACTOR, maxiter $DYN_MAXITER, min step $DYN_STEP_MIN"
echo "  second run    ${RUN_YS:+from $RUN_YS}${RUN_YS:-none}"
echo "  chunks        ${DYN_CHUNKS:-$DYN_CHUNKS_ALL} of $RUN_CHUNK_SIZE, $RUN_CPUS cpus"
echo "  oggm          $OGGM_REPO@$OGGM_COMMIT${OGGM_PATCH:+ + $(basename $OGGM_PATCH)}"
[ -n "$AFTER" ] && echo "  after job     $AFTER"
echo "------------------------------------------------------------------"

for RUN in "$@"; do
    parse_run "$RUN"
    EXP_DIR=$V7_EXPS_DIR/$EXP
    RUN_OUT=$OUT_DIR/$RUN
    mkdir -p "$LOGDIR/$RUN"
    dep=""
    [ -n "$AFTER" ] && dep="--dependency=afterok:$AFTER --kill-on-invalid-dep=yes"
    line="$RUN (gdirs $EXP, shift '${SHIFT}', ys ${YS_RUN:-$DYN_YS}, calib ${CALIB_RUN:-$CALIB}):"
    if [ "$ONLY" != "merge" ]; then
        jid=$(DYN_YS=${YS_RUN:-$DYN_YS} sbatch --parsable --export=ALL,DYN_YS=${YS_RUN:-$DYN_YS},CALIB=${CALIB_RUN:-$CALIB} --job-name="r5_${RUN}" \
                     --array=${DYN_CHUNKS:-$DYN_CHUNKS_ALL} \
                     --cpus-per-task=$RUN_CPUS --time=$DYN_RUN_TIME \
                     --mail-user=$MAIL $dep \
                     --output="$LOGDIR/$RUN/dyn_%A_%a.out" \
                     "$HERE/run_dyn_slurm.sh" "$EXP_DIR" "$RUN_OUT/chunks" "$TEST" "$SHIFT")
        dep="--dependency=afterok:$jid --kill-on-invalid-dep=yes"
        line="$line  dyn=$jid"
    fi
    MERGE="$LOCAL_PY $HERE/merge_chunks.py $RUN_OUT/chunks $RUN_OUT/${DYN_ID}_reg11.nc $DYN_ID"
    [ -n "$RUN_YS" ] && MERGE="$MERGE && $LOCAL_PY $HERE/merge_chunks.py $RUN_OUT/chunks $RUN_OUT/${RUN_ID}_reg11.nc $RUN_ID"
    jid=$(sbatch --parsable --job-name="r5_${RUN}_merge" \
                 --ntasks=1 --cpus-per-task=2 --time=00:30:00 \
                 --mail-user=$MAIL --mail-type=FAIL $dep \
                 --output="$LOGDIR/$RUN/merge_%j.out" \
                 --wrap="$MERGE")
    echo "$line  merge=$jid"
done