#!/bin/bash
# Submit a round 7 ensemble run: all HIST and NAT members from the AdaPBS L4
# gdirs on v8 (one calibration, on HIST ens010, for all members).
#
#   ./submit_ens.sh <dataset name> [--test]
#   e.g. ./submit_ens.sh ens_v8_adapbs
#
#   --test        chunk 8 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)
#   AFTER=jids    start after these jobs ended OK (e.g. the forcing builds)
#
# Forcing (common.sh; data/build_laprec_precip.py, data/build_laprec_nat.py):
#   HIST  T  Mira's BE_on_all strand (own LMR < 1850, shared after)
#         P  LAPrec-resampled own < 1850, shared resampled 1850-1899,
#            LAPrec 1900-2023, ERA5 mapped 2024-2025
#   NAT   T  lmr_seas_fullsplit_oldBE_newnat
#         P  LAPrec-resampled from its own NAT temperature (== HIST < 1850)
# Members: lmr_BE_on_all_laprec_v20261031/members_{hist,nat}.txt (192 / 197).
#
# Output: output_ens/<name>/lmr_mira_{hist,nat}_ens###_reg11.nc,
# lmr_mira_hist_reg11.nc -> ens010 (the calibration member), members.txt.
set -e

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

if [ -z "$1" ]; then
    echo "Usage: $0 <dataset name> [--test]" >&2
    exit 1
fi
NAME=$1; TEST=$2
export ENS_YS=${ENS_YS:-1500}
export ENS_BATCH_SIZE=20 ENS_N_CHUNKS=$RUN_N_CHUNKS
ENS_TIME=12:00:00
M=$DATA_DIR/$LAPREC

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
    { sed 's/^/hist /' $M/members_hist.txt; sed 's/^/nat /' $M/members_nat.txt; } > "$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:-8}; 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         $TREE_URL (L4)"
echo "  HIST          $HIST_URL"
echo "  NAT           P $NAT_URL, T $NAT_T_URL"
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${AFTER:+, after $AFTER}"
echo "  oggm          $OGGM_REPO@$OGGM_COMMIT"
echo "------------------------------------------------------------------"

jid=$(sbatch --parsable --job-name="r7e_${NAME}" --array=$ARRAY \
             ${AFTER:+--dependency=afterok:$AFTER} \
             --cpus-per-task=$RUN_CPUS --time=$ENS_TIME \
             --mail-user=$MAIL --export=ALL \
             --output="$LOGDIR/run_%A_%a.out" \
             "$HERE/run_ens_slurm.sh" "$OUT/chunks")
echo "$NAME: run=$jid"
echo "merge on the login node when done (the slurm merge hangs on NFS imports):"
echo "  MERGE_JOBS=8 MAIN_MEMBER=010 $LOCAL_PY $HERE/merge_ens.py $OUT $N_EXPECT"