#!/bin/bash
# Submit the round 7 permutation runs: every 1850 state of an L5 state tree x every NAT
# climate, from 1850 (conditional NAT ensembles; see README.md).
#
#   ./submit_perm.sh <state tree> <dataset name> [--test]
#   e.g. ./submit_perm.sh a5v10_1850 perm_a5v10
#
#   --test        chunk 8, 4 glaciers, NAT 023 and 118 only, from the test state tree,
#                 into test_run/<name>/
#   MEMBERS=file  NAT members (one number per line) instead of all 197
#   TASKS=ids     resubmit these array tasks only
#   AFTER=jid     start after job jid ended OK (e.g. the state tree build)
#
# Then (one slurm job, afterok): merge (round7/merge_ens.py), the HIST strands of the
# states (collect_hist.py), the regional CSVs (preprocess_perm.py).
# Output: output/<name>/perm_state###_{nat_ens###,hist}_reg11.nc, filtered/<name>/.
set -e

HERE=$(cd "$(dirname "$0")" && pwd)
source "$HERE/../common.sh"
[ -n "$2" ] || { echo "Usage: $0 <state tree> <dataset name> [--test]" >&2; exit 1; }
TREE=$1; NAME=$2; TEST=$3
export R6_YS=$PERM_YS R6_STATES="${STATES:-$PERM_STATES}"
export R6_NAT_T_URL=$NAT_T_URL R6_P_URL=$NAT_URL R6_BATCH_SIZE=$PERM_BATCH_SIZE
M=$DATA_DIR/$LAPREC/members_nat.txt
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
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 "023\n118\n" > "$MEMBERS_FILE"
else
    cp $M "$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, LAPrec-resampled P), $N_BATCHES batches"
echo "  runs          $(( N_MEMBERS * N_STATES ))"
echo "  tasks         $ARRAY, $PERM_CPUS cpus, $TIME${AFTER:+, after $AFTER}"
echo "------------------------------------------------------------------"
jid=$(sbatch --parsable --job-name="r7p_${NAME}" --array=$ARRAY \
             ${AFTER:+--dependency=afterok:$AFTER --kill-on-invalid-dep=yes} \
             --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")
POST="MERGE_JOBS=8 MAIN_MEMBER=none $LOCAL_PY -u $R7_DIR/merge_ens.py $OUT $N_EXPECT && \
$LOCAL_PY -u $HERE/collect_hist.py $OUT $TREE_DIR && \
N_JOBS=8 $LOCAL_PY -u $HERE/preprocess_perm.py $NAME"
[ "$TEST" = "--test" ] && POST="${POST%%&& N_JOBS*}"
mid=$(sbatch --parsable --job-name="r7p_${NAME}_post" --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/post_%j.out" --wrap="$POST")
echo "$NAME: run=$jid  post=$mid"