#!/bin/bash
# Submit a round 4 ensemble run: all the HIST and NAT members of
# lmr_seas_fullsplit_oldBE on one v5 gdir set, with the calibration made with
# ens183 HIST (only the forcing changes, as in round3).
#
#   HIST_DIR=jjas_hists_v20261005 \
#   MELT_F_CSV=sensitivity/output/lam8_jjas_dyn/dyn_calib_reg11.csv \
#       ./submit_ens.sh ensemble lam8_jjas lam8_jjas_dyn
#
#   $1   study folder (e.g. ensemble)
#   $2   v5 gdir experiment (lmr_gdirs_v5_202610/exps/<exp>)
#   $3   dataset name -> <study>/output/<name>/
#   $4   --test: chunk 39 x (hist 183, nat 001) only, into <study>/test_run/
#
#   HIST_DIR     required: the HIST temperature folder in the unstacked round
#                (jjas_hists_v20261005, monthly_hists_v20261005, ...). Use
#                the one the gdirs were built with.
#   MELT_F_CSV   optional: the dyn_calib_reg11.csv of a DYN_CALIB run, whose
#                melt_f is put back into the gdirs (empty: static melt_f)
#   AFTER        optional: a job id the runs wait for (afterok). The gdirs
#                and MELT_F_CSV need not exist yet then.
#   NAT_ROUND    optional: the unstacked round with the NAT temperatures
#                (ensemble/tp_nat_monthly_ens###.nc), default
#                lmr_seas_fullsplit_oldBE. Precip and elevation always come
#                from lmr_seas_fullsplit_oldBE (e.g. NAT_ROUND=
#                lmr_seas_fullsplit_oldBE_newnat, delivered 2026-10-16
#                without precip)
#   MEMBERS      optional: a members file ("hist 053" / "nat 001" per line)
#                instead of all of them
#
# Output: <study>/output/<name>/lmr_mira_{hist,nat}_ens###_reg11.nc (round3's
# names), lmr_mira_hist_reg11.nc -> the ens183 HIST member, and members.txt.
# Round3's post-processing then works as is:
#   cd <study> && python ../preprocess_mira.py <name>
#
# One array task = one chunk x one batch of ENS_BATCH_SIZE members, then one
# merge job (merge_ens.py). Recovery when tasks failed (the merge is then
# cancelled): resubmit the failed task ids with TASKS=..., then
#   $LOCAL_PY merge_ens.py <study>/output/<name> 40
set -e

HERE=$(cd "$(dirname "$0")" && pwd)
# Keep the caller's HIST_DIR: common.sh sources the v5 common.sh, which sets
# its own default HIST_DIR (that silently overrode ours once, on 2026-10-05)
HIST_DIR_ARG=$HIST_DIR
source "$HERE/common.sh"
HIST_DIR=$HIST_DIR_ARG

if [ -z "$3" ]; then
    echo "Usage: $0 <study> <v5 gdir exp> <dataset name> [--test]" >&2
    exit 1
fi
STUDY_DIR=$(mkdir -p "$1" && cd "$1" && pwd)
EXP=$2
NAME=$3
TEST=$4

DATA_DIR=/home/www/fmaussion/runs/lmr_2025/data/unstacked/lmr_seas_fullsplit_oldBE
NAT_ROUND=${NAT_ROUND:-lmr_seas_fullsplit_oldBE}
NAT_DATA_DIR=$(dirname $DATA_DIR)/$NAT_ROUND
NAT_URL=$(dirname $LMR_URL)/$NAT_ROUND
if [ ! -d "$NAT_DATA_DIR/ensemble" ]; then
    echo "ERROR: no $NAT_DATA_DIR/ensemble" >&2; exit 1
fi
if [ -z "$HIST_DIR" ] || [ ! -d "$DATA_DIR/$HIST_DIR/ensemble" ]; then
    echo "ERROR: set HIST_DIR to a folder of $DATA_DIR (got '$HIST_DIR')" >&2; exit 1
fi
if [ ! -f "$V5_EXPS_DIR/$EXP/config.sh" ]; then
    echo "ERROR: no v5 experiment '$EXP'" >&2; exit 1
fi
if [ -z "$AFTER" ] && [ ! -d "$V5_EXPS_DIR/$EXP/RGI62/b_160/L3" ]; then
    echo "ERROR: no v5 gdirs for '$EXP'" >&2; exit 1
fi
if [ -n "$MELT_F_CSV" ]; then
    MELT_F_CSV=$(realpath -m "$MELT_F_CSV")
    if [ -z "$AFTER" ] && [ ! -f "$MELT_F_CSV" ]; then
        echo "ERROR: no $MELT_F_CSV" >&2; exit 1
    fi
fi
DEP=""
[ -n "$AFTER" ] && DEP="--dependency=afterok:$AFTER --kill-on-invalid-dep=yes"

ENS_BATCH_SIZE=20
ENS_N_CHUNKS=40          # region 11, chunks of 100 (see submit.sh)
ENS_TIME=08:00:00

if [ "$TEST" = "--test" ]; then
    OUT=$STUDY_DIR/test_run/$NAME
    LOGDIR=$STUDY_DIR/_logs_test
    ENS_TIME=01:00:00
else
    OUT=$STUDY_DIR/output/$NAME
    LOGDIR=$STUDY_DIR/_logs
fi
mkdir -p "$OUT" "$LOGDIR"

# The members. A subset (MEMBERS) gets its own file, so that members.txt
# keeps listing the whole dataset
MEMBERS_FILE=$OUT/members.txt
if [ -n "$MEMBERS" ]; then
    MEMBERS_FILE=$OUT/members_rerun_$(date +%Y%m%d_%H%M%S).txt
    cp "$MEMBERS" "$MEMBERS_FILE"
elif [ "$TEST" = "--test" ]; then
    printf "hist 183\nnat 001\n" > "$MEMBERS_FILE"
else
    {
        ls $DATA_DIR/$HIST_DIR/ensemble/ | sed -n 's/^tp_monthly_ens\([0-9]*\)\.nc$/hist \1/p'
        ls $NAT_DATA_DIR/ensemble/ | sed -n 's/^tp_nat_monthly_ens\([0-9]*\)\.nc$/nat \1/p'
    } > "$MEMBERS_FILE"
fi
N_MEMBERS=$(grep -c . "$MEMBERS_FILE")
N_BATCHES=$(( (N_MEMBERS + ENS_BATCH_SIZE - 1) / ENS_BATCH_SIZE ))
if [ "$TEST" = "--test" ]; then
    ARRAY=${TASKS:-39}
else
    ARRAY=${TASKS:-0-$(( N_BATCHES * ENS_N_CHUNKS - 1 ))}
fi

echo "------------------------------------------------------------------"
echo "  dataset       $OUT"
echo "  gdirs         v5 $EXP"
echo "  HIST temp     $HIST_DIR"
echo "  NAT temp      $NAT_ROUND"
echo "  melt_f        ${MELT_F_CSV:-static (from the gdirs)}"
echo "  members       $N_MEMBERS ($(grep -c '^hist' $MEMBERS_FILE) hist, $(grep -c '^nat' $MEMBERS_FILE) nat), $N_BATCHES batches of $ENS_BATCH_SIZE"
echo "  tasks         $ARRAY ($ENS_N_CHUNKS chunks x $N_BATCHES batches), $RUN_CPUS cpus, $ENS_TIME"
echo "  oggm commit   $OGGM_COMMIT${OGGM_PATCH:+ + $(basename $OGGM_PATCH)}"
[ -n "$AFTER" ] && echo "  after job     $AFTER"
echo "------------------------------------------------------------------"

export ENS_MEMBERS=$MEMBERS_FILE ENS_BATCH_SIZE ENS_N_CHUNKS ENS_MELT_F_CSV=$MELT_F_CSV ENS_HIST_DIR=$HIST_DIR ENS_NAT_URL=$NAT_URL
jid=$(sbatch --parsable --job-name="r4e_${NAME}" \
             --array=$ARRAY \
             --cpus-per-task=$RUN_CPUS --time=$ENS_TIME \
             --mail-user=$MAIL --export=ALL $DEP \
             --output="$LOGDIR/${NAME}_run_%A_%a.out" \
             "$HERE/run_slurm.sh" $EXP "$OUT/chunks" ens)
N_EXPECT=$ENS_N_CHUNKS
[ "$TEST" = "--test" ] && N_EXPECT=1
mid=$(sbatch --parsable --job-name="r4e_${NAME}_merge" \
             --ntasks=1 --cpus-per-task=8 --time=12:00:00 \
             --mail-user=$MAIL --mail-type=FAIL \
             --dependency=afterok:$jid --kill-on-invalid-dep=yes \
             --output="$LOGDIR/${NAME}_merge_%j.out" \
             --wrap="MERGE_JOBS=8 $LOCAL_PY $HERE/merge_ens.py $OUT $N_EXPECT")
echo "$NAME: run=$jid  merge=$mid"