#!/bin/bash
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --time=08:00:00
#SBATCH --qos=normal
#SBATCH --mail-type=FAIL
#
# AdaPBS calibrations of a list of glaciers, as a slurm array: task i does the
# glaciers i, i + N, i + 2N... of the list (N = array size), one after the
# other, each with an 8-worker particle pool (calib_aletsch.py + analyse.py).
# Results are copied back after each glacier, so a killed task keeps its
# finished glaciers.
#
#   sbatch --array=0-39 --job-name=ada_<tag> run_batch_slurm.sh <id_file> <tag> [calib options]
#   ADA_TREE / ADA_EXP select the gdirs (default v7 sw_pf3.5_mf5_lam12).
#   Glaciers whose out_<tag>/<RGI>_<tag>.nc exists are skipped (resubmit to resume).
#   --slim: region runs (no pickle, no plots, float32 series from 1800).
#
# Output: adapbs/out_<tag>/, logs in adapbs/_logs/
set -e
IDS=$1
TAG=$2
shift 2
ADA=/home/www/fmaussion/runs/lmr_2025/adapbs
OGGM_REPO=https://github.com/fmaussion/oggm.git
OGGM_COMMIT=50c97b88669598c0822abd562041e32dd375dc66      # branch adapbs (incl. upstream speed-ups #2024, #2025)
CONTAINER=/home/users/fmaussion/images/oggm_v163.sif
SAMPLE_DATA_DIR=/home/www/fmaussion/runs/lmr_2025/lmr_gdirs_v5_202610/_sample_data
TASK=${SLURM_ARRAY_TASK_ID:-0}
NTASK=${SLURM_ARRAY_TASK_COUNT:-1}
MINE=$(awk -v t=$TASK -v n=$NTASK 'NF && (NR-1) % n == t' $IDS | tr '\n' ' ')
echo "task $TASK of $NTASK: $MINE"

JOB=/work/$SLURM_JOB_USER/${SLURM_ARRAY_JOB_ID:-$SLURM_JOB_ID}_$TASK
RUN=$JOB/run
mkdir -p "$RUN/dl" $ADA/out_$TAG
cp $ADA/laptop/calib_aletsch.py $ADA/laptop/analyse.py "$RUN/"
cp $ADA/aletsch_forward.py "$RUN/dl/"     # no fastdate.py: OGGM has it since #2024
ADA_EXP=${ADA_EXP:-sw_pf3.5_mf5_lam12}
ADA_TREE=${ADA_TREE:-lmr_gdirs_v7_202610}
echo "gdirs: $ADA_TREE/exps/$ADA_EXP, OGGM $OGGM_COMMIT"
# --slim (region runs): no analyse.py plots
SLIM=0; case " $* " in *" --slim "*) SLIM=1 ;; esac
# Read the gdirs from disk, not over HTTP: our trees through a cache link,
# OGGM's base L2 trees from the read-only download cache (as the gdir builds)
mkdir -p "$RUN/wd/cache/cluster.klima.uni-bremen.de" "$JOB/oggm_tmp"
ln -sfn /home/www/fmaussion "$RUN/wd/cache/cluster.klima.uni-bremen.de/~fmaussion"

singularity exec $CONTAINER bash -s <<EOF
  set -e
  export HOME="$JOB/fake_home"
  mkdir -p "\$HOME/.salem_cache" "\$HOME/.oggm"
  ln -s $SAMPLE_DATA_DIR/salem-sample-data-* "\$HOME/.salem_cache/"
  ln -s $SAMPLE_DATA_DIR/oggm-sample-data-* "\$HOME/.oggm/"
  python3 -m venv --system-site-packages "$JOB/venv"
  source "$JOB/venv/bin/activate"
  pip install --quiet --upgrade pip
  pip install --quiet pyarrow
  pip install --quiet --force-reinstall --no-deps netCDF4
  pip install --quiet --no-deps "git+$OGGM_REPO@$OGGM_COMMIT"
  cd "$RUN"
  export ADA_EXP=$ADA_EXP ADA_TREE=$ADA_TREE
  export OGGM_DOWNLOAD_CACHE=/home/data/download OGGM_DOWNLOAD_CACHE_RO=1
  export OGGM_EXTRACT_DIR="$JOB/oggm_tmp"
  mkdir -p out
  for RGI in $MINE; do
    if [ -f "$ADA/out_$TAG/\${RGI}_$TAG.nc" ]; then
      echo "=== \$RGI already done, skipped"; continue
    fi
    echo "=== \$RGI start \$(date +%T)"
    export RGI_ID=\$RGI
    t0=\$(date +%s)
    if python calib_aletsch.py --workers $SLURM_CPUS_PER_TASK --tag $TAG $* ; then
      if [ "$SLIM" = "0" ]; then python analyse.py $TAG || echo "analyse failed for \$RGI"; fi
    else
      echo "CALIB FAILED for \$RGI"
    fi
    echo "=== \$RGI done in \$(( \$(date +%s) - t0 )) s"
    # this glacier's files only (the .nc last, so that it marks completion)
    for f in out/\${RGI}_*; do
      [ -e "\$f" ] || continue
      case "\$f" in *.nc) ;; *) mv "\$f" $ADA/out_$TAG/ ;; esac
    done
    if [ -f "out/\${RGI}_$TAG.nc" ]; then mv "out/\${RGI}_$TAG.nc" $ADA/out_$TAG/; fi
  done
EOF
echo "SLURM DONE"