#!/bin/bash --login
#SBATCH --job-name=alphafold
#SBATCH --output=af_output/%x.d%j/%x.o%j
#SBATCH --error=af_output/%x.d%j/%x.e%j
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --time=0-02:00:00
#SBATCH --mem=10G
#SBATCH --partition=gpu
#SBATCH --qos=gpu
#SBATCH --gres=gpu:a100_40g:1
##SBATCH --account=<account>

shopt -q login_shell || exit 1
test -n "$SLURM_NODELIST" || exit 1
test $SLURM_NNODES -eq 1 || exit 1

module load apptainer
module load alphafold/3.0.1

test -n "$ALPHAFOLD_CONTAINER" || exit 1

set -x

ALPHAFOLD_JSON_INPUT_FILE='fold_input.json'
ALPHAFOLD_INPUT_DIR="$PWD/af_input"
ALPHAFOLD_OUTPUT_DIR="$PWD/af_output/${SLURM_JOB_NAME}.d${SLURM_JOB_ID}"

mkdir -p "$ALPHAFOLD_OUTPUT_DIR"

apptainer exec \
    --nv \
    --bind "$ALPHAFOLD_INPUT_DIR:/root/af_input" \
    --bind "$ALPHAFOLD_OUTPUT_DIR:/root/af_output" \
    "$ALPHAFOLD_CONTAINER" \
    python /app/alphafold/run_alphafold.py \
    --json_path="/root/af_input/$ALPHAFOLD_JSON_INPUT_FILE" \
    --model_dir=/root/models \
    --db_dir=/root/public_databases \
    --db_dir=/root/public_databases_fallback \
    --output_dir=/root/af_output