#!/bin/bash
# Submit round 4a runs: per v6 experiment, the dynamic melt_f calibration +
# historical run (40 chunks of 100 glaciers), then the merge.
#
#   ./submit.sh <exp> [<exp> ...]          # e.g. lam8_be_lmr21
#   ./submit.sh --test <exp>               # the 4 TEST_IDS glaciers, from the
#                                          # exp's test_run/ gdirs
#   AFTER=<jobid> ./submit.sh <exp>        # wait for that job (the gdir build)
#   DYN_CHUNKS=3,17 ./submit.sh <exp>      # redo those chunks, then the merge
#   ONLY=merge ./submit.sh <exp>           # just the merge
#
#   output/<exp>_dyn/chunks/                    per-chunk files
#   output/<exp>_dyn/lmr_mira_hist_reg11.nc     merged run, years 1-2025
#   output/<exp>_dyn/dyn_calib_reg11.csv        per-glacier calibration table
#
# <exp> is a folder of ../lmr_gdirs_v6_202610/exps. The layout of output/ is
# the one of round4/sensitivity/output.
set -e

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

TEST=""
OUT_DIR=$R4A_DIR/output
LOGDIR=$R4A_DIR/_logs
DYN_CHUNKS_ALL=0-$(( RUN_N_CHUNKS - 1 ))
if [ "$1" = "--test" ]; then
    shift
    TEST=test
    OUT_DIR=$R4A_DIR/test_run
    LOGDIR=$R4A_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] <exp> [<exp> ...]" >&2
    exit 1
fi

for EXP in "$@"; do
    EXP_DIR=$V6_EXPS_DIR/$EXP
    GDIRS=$EXP_DIR/RGI62/b_160/L3
    [ -n "$TEST" ] && GDIRS=$EXP_DIR/test_run/RGI62/b_160/L3
    if [ ! -f "$EXP_DIR/config.sh" ]; then
        echo "ERROR: no v6 experiment '$EXP'" >&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 "  experiments   $*"
echo "  output        $OUT_DIR/<exp>_dyn/"
echo "  dyn           from $DYN_YS, Hugonnet err x $DYN_ERR_FACTOR, maxiter $DYN_MAXITER, min step $DYN_STEP_MIN; run from $RUN_YS"
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 EXP in "$@"; do
    EXP_DIR=$V6_EXPS_DIR/$EXP
    # JOB_TAG of the experiment, in a subshell: config.sh re-sources the v6
    # common.sh
    TAG=$(source "$EXP_DIR/config.sh"; echo "$JOB_TAG")
    DYN_OUT=$OUT_DIR/${EXP}_dyn
    mkdir -p "$LOGDIR/$EXP"
    dep=""
    [ -n "$AFTER" ] && dep="--dependency=afterok:$AFTER --kill-on-invalid-dep=yes"
    line="${EXP}_dyn:"
    if [ "$ONLY" != "merge" ]; then
        jid=$(sbatch --parsable --job-name="r4a_${TAG}_dyn" \
                     --array=${DYN_CHUNKS:-$DYN_CHUNKS_ALL} \
                     --cpus-per-task=$RUN_CPUS --time=$DYN_RUN_TIME \
                     --mail-user=$MAIL $dep \
                     --output="$LOGDIR/$EXP/dyn_%A_%a.out" \
                     "$HERE/run_dyn_slurm.sh" "$EXP_DIR" "$DYN_OUT/chunks" "$TEST")
        dep="--dependency=afterok:$jid --kill-on-invalid-dep=yes"
        line="$line  dyn=$jid"
    fi
    jid=$(sbatch --parsable --job-name="r4a_${TAG}_merge" \
                 --ntasks=1 --cpus-per-task=2 --time=00:30:00 \
                 --mail-user=$MAIL --mail-type=FAIL $dep \
                 --output="$LOGDIR/$EXP/merge_%j.out" \
                 --wrap="$LOCAL_PY $HERE/merge_chunks.py $DYN_OUT/chunks $DYN_OUT/${RUN_ID}_reg11.nc")
    echo "$line  merge=$jid"
done