#!/bin/bash -x 

#SBATCH --account=<group>
#SBATCH --partition=std
#SBATCH --job-name=Gromacs
#SBATCH --output=slurm-%x.%N.%j.out 
#SBATCH --error=slurm-%x.%N.%j.err 

# Resource configuration
#SBATCH --nodes=2                     # Number of nodes to use for the simulation (to be adjusted)
#SBATCH --cpus-per-task=4             # Number of threads per MPI process (between 4 and 6)
#SBATCH --time=0-04:00:00             # Maximum execution time
#SBATCH --exclude=cnd[01-04,06-12]    # Exclude IVY nodes incompatible with GROMACS

# Display SLURM environment variables
env | grep -i slurm

# Change directory to the submission directory
cd "$SLURM_SUBMIT_DIR"

# Create a temporary directory for the job
WORKDIR="$SCRATCHDIR/job.$SLURM_JOB_ID.$USER"
mkdir -p "$WORKDIR"

# Load necessary modules for GROMACS
module purge
module load gromacs/gcc-11.2.0/2022.6

# Source the GROMACS environment file to set up necessary variables
source "$GMXROOT/bin/GMXRC"

# Calculate the number of tasks per node - 4 corresponds to cpus-per-task
export NTASKS_PER_NODE=$((SLURM_CPUS_ON_NODE / 4)) 

# Set up environment variables for parallelism
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
ulimit -s unlimited  # Increase stack size to avoid memory errors

# Copy submission files to the working directory
cp "$SLURM_SUBMIT_DIR"/* "$WORKDIR"

# Change directory to the working directory
cd "$WORKDIR"

# Log the start date and time of the computation
echo "Computation start: $(date)" >> "$SLURM_SUBMIT_DIR/slurm-${SLURM_JOB_ID}.log"

# Run the simulation with mdrun
srun --ntasks-per-node=$NTASKS_PER_NODE gmx_mpi mdrun [arguments]

#--- End of GROMACS execution -------------------------------------------------------------

# Log the end date and time of the computation
echo "Computation end: $(date)" >> "$SLURM_SUBMIT_DIR/slurm-${SLURM_JOB_ID}.log"

# Create a directory for output files
OUTDIR="$SLURM_SUBMIT_DIR/outdir.$SLURM_JOB_ID"  # Define the output directory path
mkdir -p "$OUTDIR"  # Create the output directory (if it does not already exist)

# Move all working files to the output directory
cp "$WORKDIR"/* "$OUTDIR"

# Optional: uncomment the following line to remove the working directory
# rm -rf "$WORKDIR"