#!/bin/bash
# Submit a round 5 ensemble run: HIST and NAT members on one v7 gdir set, with
# the melt_f of its round5 calibration (one calibration for all members).
#
#   ./submit_ens.sh <v7 exp> <calibration run> <dataset name> [--test]
#   e.g. ./submit_ens.sh sw_pf3.5_mf5_lam12 sw_pf3.5_mf5_lam12_ys1500 ens_pf3.5_mf5_lam12
#
#   --test        chunk 27 x (hist 010, nat 010) only, into test_run/<name>/
#   MEMBERS=file  a members file ("hist 010" / "nat 001" per line) instead of
#                 all of them
#   TASKS=ids     resubmit these array tasks only (recovery)
#
# Forcing (2026-10-25, see CLAUDE.md "Plan for the talk"):
#   HIST  T  lmr_seas_fullsplit_oldBE_BE_on_all  (Mira's strands, shared from 1850)
#         P  own before 1850, ens$ENS_HIST_P_SHARED's from 1850
#   NAT   T  lmr_seas_fullsplit_oldBE_newnat (the v2 NAT of 2026-10-17)
#         P  own
#   P and elev from lmr_seas_fullsplit_oldBE
# HIST members: those with both a strand and a precip file; NAT: all newnat.
#
# Output: output_ens/<name>/lmr_mira_{hist,nat}_ens###_reg11.nc,
# lmr_mira_hist_reg11.nc -> the calibration member (ens010), members.txt.
set -e

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

if [ -z "$3" ]; then
    echo "Usage: $0 <v7 exp> <calibration run> <dataset name> [--test]" >&2
    exit 1
fi
EXP=$1; CALRUN=$2; NAME=$3; TEST=$4
EXP_DIR=$V7_EXPS_DIR/$EXP
BDIR=$(source "$EXP_DIR/config.sh"; printf "b_%03d" "$BORDER")
[ -d "$EXP_DIR/RGI62/$BDIR/L3" ] || { echo "ERROR: no gdirs $EXP_DIR/RGI62/$BDIR/L3" >&2; exit 1; }
MELT_F_CSV=$HERE/output/$CALRUN/dyn_calib_reg11.csv
[ -f "$MELT_F_CSV" ] || { echo "ERROR: no $MELT_F_CSV" >&2; exit 1; }

U=/home/www/fmaussion/runs/lmr_2025/data/unstacked
HIST_T=lmr_seas_fullsplit_oldBE_BE_on_all
NAT_T=lmr_seas_fullsplit_oldBE_newnat
PR=lmr_seas_fullsplit_oldBE
export ENS_HIST_P_SHARED=${ENS_HIST_P_SHARED:-148}
export ENS_HIST_P_FROM=1850
export ENS_YS=${ENS_YS:-1000}
export ENS_HIST_T_URL=$DATA_URL/$HIST_T ENS_NAT_T_URL=$DATA_URL/$NAT_T ENS_P_URL=$DATA_URL/$PR
export ENS_MELT_F_CSV=$MELT_F_CSV
export ENS_BATCH_SIZE=20 ENS_N_CHUNKS=$RUN_N_CHUNKS
ENS_TIME=12:00:00

if [ "$TEST" = "--test" ]; then
    OUT=$HERE/test_run/$NAME; LOGDIR=$HERE/_logs_test/ens_$NAME; ENS_TIME=02:00:00
else
    OUT=$HERE/output_ens/$NAME; LOGDIR=$HERE/_logs/ens_$NAME
fi
mkdir -p "$OUT" "$LOGDIR"

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 010\nnat 010\n" > "$MEMBERS_FILE"
else
    {
        for f in $U/$HIST_T/ensemble/tp_monthly_ens*.nc; do
            m=$(basename $f .nc); m=${m#tp_monthly_ens}
            [ -e $U/$PR/ensemble/p_monthly_ens$m.nc ] && echo "hist $m"
        done
        ls $U/$NAT_T/ensemble/ | sed -n 's/^tp_nat_monthly_ens\([0-9]*\)\.nc$/nat \1/p'
    } > "$MEMBERS_FILE"
fi
export ENS_MEMBERS=$MEMBERS_FILE
N_MEMBERS=$(grep -c . "$MEMBERS_FILE")
N_BATCHES=$(( (N_MEMBERS + ENS_BATCH_SIZE - 1) / ENS_BATCH_SIZE ))
if [ "$TEST" = "--test" ]; then
    ARRAY=${TASKS:-27}
    N_EXPECT=1
else
    ARRAY=${TASKS:-0-$(( N_BATCHES * ENS_N_CHUNKS - 1 ))}
    N_EXPECT=$ENS_N_CHUNKS
fi

echo "------------------------------------------------------------------"
echo "  dataset       $OUT"
echo "  gdirs         v7 $EXP ($BDIR), melt_f from $CALRUN"
echo "  HIST          T $HIST_T, P own < $ENS_HIST_P_FROM then ens$ENS_HIST_P_SHARED"
echo "  NAT           T $NAT_T, P own"
echo "  runs from     $ENS_YS"
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, $RUN_CPUS cpus, $ENS_TIME"
echo "  oggm          $OGGM_REPO@$OGGM_COMMIT"
echo "------------------------------------------------------------------"

jid=$(sbatch --parsable --job-name="r5e_${NAME}" --array=$ARRAY \
             --cpus-per-task=$RUN_CPUS --time=$ENS_TIME \
             --mail-user=$MAIL --export=ALL \
             --output="$LOGDIR/run_%A_%a.out" \
             "$HERE/run_ens_slurm.sh" "$EXP_DIR" "$OUT/chunks")
mid=$(sbatch --parsable --job-name="r5e_${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/merge_%j.out" \
             --wrap="MERGE_JOBS=8 MAIN_MEMBER=010 $LOCAL_PY $HERE/merge_ens.py $OUT $N_EXPECT")
echo "$NAME: run=$jid  merge=$mid"