#!/bin/bash

#####
##### Disclaimer: this script can be used only on SLURM type systems.
##### For PBS systems `job_script.sh` needs to be corrected to use the `qsub` syntax
#####

# Before running, make sure that:
#
# 1) The julia environmental variables (defined in this file) point to the right path
#
# 2) the NGPUS_PER_NODE variable is correct (in this file)
#
# 3) If the system is equipped with nsys profiler it is possible to enable a trace with PROFILE_TRACE=1 (in this file)
#
# 4) Oceananigans is instantiated with the correct MPI build:
# 	(run these lines in a gpu node substituting modules and paths)
# 	$ module load my_cuda_module
# 	$ module load my_cuda_aware_mpi_module
# 	$ export JULIA_DEPOT_PATH="/path/to/depot"
# 	$ export JULIA="/path/to/julia"
#   $ $JULIA --check-bounds=no -e 'using Pkg; Pkg.add("MPIPreferences");'
# 	$ $JULIA --project --check-bounds=no -e 'using MPIPreferences; MPIPreferences.use_system_binaries()'
#   $ $JULIA --project --check-bounds=no -e 'using Pkg; Pkg.build("MPI")'
#   $ $JULIA --project --check-bounds=no -e 'using Pkg; Pkg.instantiate()'
#
# 5) correct modules are loaded in job_script.sh
#
# 6) SBATCH variables in the job_script.sh file are correct (check memory, time)
#
# 7) The system has at least max(RX) * max(RY) gpus
#
# 8) Choose if testing the hydrostatic or nonhydrostatic model (SIMULATION variable in this file)
#
# 9) Choose if measuring the weak or strong scaling (SCALING variable in this file)
#
# Finally -> $ ./run_tests.sh

# Julia specific enviromental variables
export JULIA_DEPOT_PATH="/path/to/depot"
export JULIA="/path/to/julia"

# PROFILE_TRACE=1 only if the system is equipped with nsys
export PROFILE_TRACE=0

# Number of gpus per node
export NGPUS_PER_NODE=4

# Choice between nonhydrostatic and hydrostatic
export SIMULATION=nonhydrostatic
# Choice between strong and weak
export SCALING=weak

for RX in 1 2 4 8 16 32 64; do
    for RY in 1 2 4 8 16 32 64; do

		export RX
        export RY

		if test $SIMULATION = "hydrostatic"; then
			if test $SCALING = "weak"; then
				# Grid size for Weak scaling tests (Hydrostatic)
				export NX=$((1440 * RX))
				export NY=$((600 * RY))
				export NZ=100
			else
				# Grid size for Strong scaling tests (Hydrostatic)
				export NX=1440
				export NY=600
				export NZ=100
			fi
		else
			if test $SCALING = "weak"; then
				# Grid size for Weak scaling tests (Nonhydrostatic)
				export NX=$((512 * RX))
				export NY=$((512 * RY))
				export NZ=256
			else
				# Grid size for Strong scaling tests (Nonhydrostatic)
				export NX=512
				export NY=512
				export NZ=256
			fi
		fi

		RANKS=$((RX * RY))

		export NNODES=$((RANKS / NGPUS_PER_NODE))
		export NTASKS=$NGPUS_PER_NODE

		echo ""
		echo "(RX, RY) = $RX, $RY"
		echo "(NX, NY) = $NX, $NY"
		echo "(NNODES, NTASKS) = $NNODES, $NTASKS"

		# ====================================================== #
		# ================== RUN SCALING TEST ================== #
		# ====================================================== #

		sbatch -N ${NNODES} --gres=gpu:${NTASKS} --ntasks-per-node=${NTASKS} job_script.sh

		# Use qsub on PBS systems!!!
		# qsub pbs_job_script.sh
    done
done
