#!/bin/bash
#SBATCH --job-name=EHD_CNN
#SBATCH --nodes=5
#SBATCH --exclusive
#SBATCH --time=08:00:00
#SBATCH --output=/home/tanpv/Progs/claude/HW/logs/cnn_%j.out
#SBATCH --error=/home/tanpv/Progs/claude/HW/logs/cnn_%j.err
set -euo pipefail
ROOT=/home/tanpv/Progs/claude/HW
PY=/home/tanpv/.conda/envs/ehd/bin/python
# torch needs the conda libstdc++ ahead of the system one (no CXXABI_1.3.9 in /lib64)
export LD_LIBRARY_PATH=/home/tanpv/.conda/envs/ehd/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
cd "$ROOT"
TAG=${TAG:-_cnn}
EPOCHS=${EPOCHS:-8}
PHASES=${PHASES:-"fit predict finalize"}
has() { [[ " $PHASES " == *" $1 "* ]]; }
export OMP_NUM_THREADS=20 MKL_NUM_THREADS=20

# The four fold models and the final model are independent, so they train at the
# same time, one node each -- same shape as the fit phase of the row pipeline.
if has fit; then
  echo "--- cnn fit (4 folds + full, one node each) ---"
  srun --nodes=5 --ntasks=5 --ntasks-per-node=1 --cpus-per-task=20 \
       --kill-on-bad-exit=1 bash -c \
       "FOLDS=(0 1 2 3 full); $PY src/train_cnn.py fit \
          --fold \${FOLDS[\$SLURM_PROCID]} --tag '$TAG' --epochs $EPOCHS --threads 20"
fi
if has predict; then
  echo "--- cnn predict ---"
  srun --nodes=1 --ntasks=1 --cpus-per-task=20 \
       $PY src/train_cnn.py predict --tag "$TAG" --threads 20
fi
# Scored by exactly the same finalize the row-based models go through.
# --neg-keep 1.0 because train_cnn.py keeps every out-of-fold row: the row
# pipeline thins its non-events to 15% and finalize undoes that reweighting, so
# leaving the default here would reweight negatives that were never thinned and
# push the threshold ~7x too high (FBI 0.02 instead of ~1).
if has finalize; then
  echo "--- finalize (shared with the row pipeline) ---"
  srun --nodes=1 --ntasks=1 --cpus-per-task=20 \
       $PY src/train_national.py finalize --model cnn --fset MAP --tag "$TAG" \
           --neg-keep 1.0
fi
echo "=== done ==="