#!/bin/bash

MODULE_NAME=nnunet_for_pytorch
MODULE_VERSION=21.11.0
WORK_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/PyTorch/Segmentation/nnUNet
CONTAINER_DIR=/proj/nsc_testing/xuan/containers/${MODULE_NAME}_${MODULE_VERSION}.sif
SBATCH_DIR=$WORK_DIR/sbatch_scripts/benchmark_${6}_${5}_dim${1}_nodes${2}_gpus${3}_batchsize_${4}.sbatch
SBATCH_OUT_DIR=$WORK_DIR/sbatch_out/benchmark_${6}_${5}_dim${1}_nodes${2}_gpus${3}_batchsize_${4}.out
LOG_DIR=benchmark_${6}_${5}_dim${1}_nodes${2}_gpus${3}_batchsize_${4}_amp.json
mkdir -p $WORK_DIR/sbatch_out $WORK_DIR/sbatch_scripts

cat <<EOT >  $SBATCH_DIR
#!/bin/bash

#SBATCH -A nsc
#SBATCH --nodes=${2}
#SBATCH --gpus=${3}
#SBATCH --time=0-0:10:00
#SBATCH --output=$SBATCH_OUT_DIR

EOT

if [ "${6}" = "thin" ]; then
    cat <<EOT >>  $SBATCH_DIR
#SBATCH -C "thin"    
#SBATCH --reservation=nodeimage
EOT
else
    cat <<EOT >>  $SBATCH_DIR
#SBATCH -C "fat"
EOT
fi

cat <<EOT >>  $SBATCH_DIR

rm -f $WORK_DIR/results/$LOG_DIR
apptainer exec --nv -B ${WORK_DIR}/data:/data -B ${WORK_DIR}/results:/results --pwd /workspace/nnunet_pyt $CONTAINER_DIR python scripts/benchmark.py --mode ${5} --gpus ${3} --dim ${1} --batch_size ${4} --amp --logname='$LOG_DIR'

EOT