#!/bin/bash
# Submit the round 6 permutation runs: every 1850 state of an L4 state tree
# x every NAT climate, from 1850 (see README.md).
#
#   ./submit_perm.sh <v7 exp> <state tree> <dataset name> [--test]
#   e.g. ./submit_perm.sh sw_pf3.5_mf5_lam12 sw_pf3.5_mf5_lam12_1850 perm_pf3.5_mf5_lam12
#
#   --test        chunk 8, 4 glaciers, NAT 010 and 041 only, from the test
#                 state tree, into test_run/<name>/
#   MEMBERS=file  NAT members (one number per line) instead of all of them
#   TASKS=ids     resubmit these array tasks only (recovery)
#
# Output: output/<name>/perm_state###_nat_ens###_reg11.nc (merged by
# round5/merge_ens.py), then perm_state###_hist_reg11.nc (collect_hist.py:
# the L5 HIST strands of the state tree), members.txt. The merge is printed
# as a command for the login node (MERGE_JOB=1: a slurm job after the runs).
set -e

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

if [ -z "$3" ]; then
    echo "Usage: $0 <v7 exp> <state tree> <dataset name> [--test]" >&2
    exit 1
fi
EXP=$1; TREE=$2; NAME=$3; TEST=$4
export R6_EXP_DIR=$V7_EXPS_DIR/$EXP
BDIR=$(source "$R6_EXP_DIR/config.sh"; printf "b_%03d" "$BORDER")
export R6_YS=$PERM_YS R6_STATES="${STATES:-$PERM_STATES}"
export R6_NAT_T_URL=$NAT_T_URL R6_P_URL=$P_URL R6_BATCH_SIZE=$PERM_BATCH_SIZE

if [ "$TEST" = "--test" ]; then
    TREE_DIR=$STATES_DIR/test_run/$TREE
    export R6_STATES_URL="$STATES_URL_BASE/test_run/$TREE/"
    export R6_TEST_IDS="RGI60-11.00897 RGI60-11.00890 RGI60-11.00846 RGI60-11.00807"
    OUT=$HERE/test_run/$NAME; LOGDIR=$HERE/_logs_test/$NAME; TIME=01:00:00
else
    TREE_DIR=$STATES_DIR/$TREE
    export R6_STATES_URL="$STATES_URL_BASE/$TREE/"
    OUT=$HERE/output/$NAME; LOGDIR=$HERE/_logs/$NAME; TIME=$PERM_TIME
fi
[ -d "$TREE_DIR/RGI62/$BDIR/L4" ] || { echo "ERROR: no L4 in $TREE_DIR" >&2; exit 1; }
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 "010\n041\n" > "$MEMBERS_FILE"
else
    ls $DATA_DIR/lmr_seas_fullsplit_oldBE_newnat/ensemble/ \
        | sed -n 's/^tp_nat_monthly_ens\([0-9]*\)\.nc$/\1/p' > "$MEMBERS_FILE"
fi
export R6_MEMBERS=$MEMBERS_FILE
N_MEMBERS=$(grep -c . "$MEMBERS_FILE")
N_BATCHES=$(( (N_MEMBERS + PERM_BATCH_SIZE - 1) / PERM_BATCH_SIZE ))
if [ "$TEST" = "--test" ]; then
    ARRAY=${TASKS:-8}; N_EXPECT=1
else
    ARRAY=${TASKS:-0-$(( N_BATCHES * PERM_N_CHUNKS - 1 ))}; N_EXPECT=$PERM_N_CHUNKS
fi
N_STATES=$(echo $R6_STATES | wc -w)

echo "------------------------------------------------------------------"
echo "  dataset       $OUT"
echo "  states        $R6_STATES (start of $R6_YS) from $R6_STATES_URL"
echo "  NAT climates  $N_MEMBERS (T newnat, own P), $N_BATCHES batches of $PERM_BATCH_SIZE"
echo "  runs          $(( N_MEMBERS * N_STATES ))"
echo "  tasks         $ARRAY, $PERM_CPUS cpus, $TIME${R6_TEST_IDS:+ (test: $R6_TEST_IDS)}"
echo "  oggm          $OGGM_REPO@$OGGM_COMMIT"
echo "------------------------------------------------------------------"

jid=$(sbatch --parsable --job-name="r6_${NAME}" --array=$ARRAY \
             --cpus-per-task=$PERM_CPUS --time=$TIME \
             --mail-user=$MAIL --export=ALL \
             --output="$LOGDIR/run_%A_%a.out" \
             "$HERE/run_perm_slurm.sh" "$OUT/chunks")
MERGE_CMD="MERGE_JOBS=8 $LOCAL_PY $R5_DIR/merge_ens.py $OUT $N_EXPECT && $LOCAL_PY $HERE/collect_hist.py $OUT $TREE_DIR"
if [ "${MERGE_JOB:-0}" = "1" ]; then
    mid=$(sbatch --parsable --job-name="r6_${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_CMD")
    echo "$NAME: run=$jid  merge=$mid"
else
    # The merge job sat > 20 min without merging a file in the test (python
    # imports over NFS on the nodes); on the login node it takes seconds.
    echo "$NAME: run=$jid"
    echo "When done, merge on the login node:"
    echo "  $MERGE_CMD"
fi