#!/bin/bash
# Submit round 4 runs: one chunked historical run per v5 experiment.
#
#   ./submit.sh sensitivity                  # all the v5 experiments
#   ./submit.sh sensitivity default mf7      # or just these
#   ./submit.sh sensitivity --test default   # smoke test: one small chunk
#
# For each experiment, this submits an array job (one task per chunk of
# $RUN_CHUNK_SIZE glaciers) and, after it, a merge job:
#
#   <study>/output/<exp>/chunks/lmr_mira_hist_reg11_chunk###.nc
#   <study>/output/<exp>/lmr_mira_hist_reg11.nc        (merged)
#
# The merged file has round3's layout and name, so round3's post-processing
# (preprocess_mira.py) and plots work on <study>/output unchanged.
#
# --test runs chunk 39 only (28 small glaciers) into <study>/test_run/.
#
# DYN_CALIB=1 first calibrates melt_f per glacier with a dynamic run against
# Hugonnet (see run_hist.py), then runs with it. The output is a separate
# experiment, <exp>_dyn, so the uncalibrated run is kept.
#
#   DYN_CALIB=1 ./submit.sh sensitivity lam8_ib     # -> output/lam8_ib_dyn
#
# AFTER=<jobid> makes the runs wait for that job, e.g. the build of the v5
# gdirs (the gdirs need not exist yet then).
#
# Recovery, for when some chunks failed (the merge job is then cancelled):
#
#   CHUNKS=3,17 ./submit.sh sensitivity mf7      # redo those, then merge
#   ONLY=merge ./submit.sh sensitivity mf7       # just the merge
set -e

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

if [ -z "$1" ]; then
    echo "Usage: $0 <study folder> [--test] [exp ...]" >&2
    exit 1
fi
STUDY_DIR=$(mkdir -p "$1" && cd "$1" && pwd)
shift

OUT_DIR=$STUDY_DIR/output
LOGDIR=$STUDY_DIR/_logs
if [ "$1" = "--test" ]; then
    shift
    OUT_DIR=$STUDY_DIR/test_run
    LOGDIR=$STUDY_DIR/_logs_test
    CHUNKS=${CHUNKS:-39}
    RUN_TIME=01:00:00
fi

EXPS="$@"
if [ -z "$EXPS" ]; then EXPS=$(ls "$V5_EXPS_DIR"); fi

# 40 chunks of 100 in region 11 (oggm_prepro_chunks)
case $RUN_CHUNK_SIZE in
    100)  N_CHUNKS=40 ;;
    1000) N_CHUNKS=4  ;;
    *) echo "no chunk count for RUN_CHUNK_SIZE=$RUN_CHUNK_SIZE" >&2; exit 1 ;;
esac
ARRAY=${CHUNKS:-0-$(( N_CHUNKS - 1 ))}

echo "------------------------------------------------------------------"
echo "  study         $STUDY_DIR"
echo "  output        $OUT_DIR/<exp>/"
echo "  experiments   $(echo $EXPS)"
echo "  run           $RUN_ID, from year $RUN_YS, chunks $ARRAY of $RUN_CHUNK_SIZE, $RUN_CPUS cpus"
echo "  oggm commit   $OGGM_COMMIT${OGGM_PATCH:+ + $(basename $OGGM_PATCH)}"
[ -n "$AFTER" ] && echo "  after job     $AFTER"
echo "------------------------------------------------------------------"

for EXP in $EXPS; do
    GDIR_DIR=$V5_EXPS_DIR/$EXP/RGI62/b_160/L3
    if [ ! -f "$V5_EXPS_DIR/$EXP/config.sh" ]; then
        echo "ERROR: no v5 experiment '$EXP'" >&2
        exit 1
    fi
    if [ -z "$AFTER" ] && [ ! -d "$GDIR_DIR" ]; then
        echo "ERROR: no v5 gdirs for '$EXP': $GDIR_DIR" >&2
        exit 1
    fi
done

DYN_ARG=""; SUFFIX=""
if [ "$DYN_CALIB" = "1" ]; then
    DYN_ARG=dyn; SUFFIX=_dyn
    RUN_TIME=$DYN_RUN_TIME
    echo "  dynamic melt_f calibration: runs from $DYN_YS, Hugonnet error x $DYN_ERR_FACTOR, maxiter $DYN_MAXITER, min step $DYN_STEP_MIN, output <exp>_dyn"
fi

mkdir -p "$LOGDIR"
for EXP in $EXPS; do
    OUT_NAME=$EXP$SUFFIX
    EXP_OUT=$OUT_DIR/$OUT_NAME
    line="$OUT_NAME:"
    dep=""
    [ -n "$AFTER" ] && dep="--dependency=afterok:$AFTER --kill-on-invalid-dep=yes"
    if [ "$ONLY" != "merge" ]; then
        jid=$(sbatch --parsable --job-name="r4_${OUT_NAME}" \
                     --array=$ARRAY \
                     --cpus-per-task=$RUN_CPUS --time=$RUN_TIME \
                     --mail-user=$MAIL $dep \
                     --output="$LOGDIR/${OUT_NAME}_run_%A_%a.out" \
                     "$HERE/run_slurm.sh" $EXP "$EXP_OUT/chunks" $DYN_ARG)
        dep="--dependency=afterok:$jid --kill-on-invalid-dep=yes"
        line="$line  run=$jid"
    fi
    jid=$(sbatch --parsable --job-name="r4_${OUT_NAME}_merge" \
                 --ntasks=1 --cpus-per-task=2 --time=00:30:00 \
                 --mail-user=$MAIL --mail-type=FAIL $dep \
                 --output="$LOGDIR/${OUT_NAME}_merge_%j.out" \
                 --wrap="$LOCAL_PY $HERE/merge_chunks.py $EXP_OUT/chunks $EXP_OUT/${RUN_ID}_reg11.nc")
    echo "$line  merge=$jid"
done