#!/bin/bash

set -e

# by default, this script runs precompiled binaries
# set USE_PRECOMPILED_BINARY to 0 to build binary from source
USE_PRECOMPILED_BINARY=1
# default path to precompiled binaries
EXEPATH=/usr/bin

# by default set 10000 dofs per process to just check that tests run successfully
DOFS_PER_PROCESS=10000

usage()
{
  echo "Usage: $0 [-s] [-d DOFS_PER_PROCESS] [-p EXEPATH]"
  exit 0
}

while getopts "h?sd:p:" opt; do
    case "$opt" in
    s)  USE_PRECOMPILED_BINARY=0
	EXEPATH=.
        ;;
    d)  DOFS_PER_PROCESS=$OPTARG
        ;;
    p)  EXEPATH=$OPTARG
        ;;
    h|?)
        usage
        ;;
    esac
done


# test parallel jobs, but don't load up more than 3
NPROC=$( nproc )
MAX_CPU=3
N_CPU=$(( NPROC > MAX_CPU ? MAX_CPU : NPROC ))

export OMPI_MCA_plm_rsh_agent=/bin/false
export OMPI_MCA_rmaps_base_oversubscribe=1

SOURCE_DIR=/usr/share/dolfinx/performance-tests/src
SOURCE_BUILD_DIR=dolfinx-performance-tests-build

# total DOFs for strong scaling determined by DOFS_PER_PROCESS
N_DOF_STRONG=$(( N_CPU * DOFS_PER_PROCESS ))
N_DOF_WEAK=$DOFS_PER_PROCESS

EXE_NAME=dolfinx-scaling-test
EXE=${EXEPATH}/${EXE_NAME}

# build test binary from source if required
if [ "x$USE_PRECOMPILED_BINARY" = "x0" ]; then
  cp -r $SOURCE_DIR $SOURCE_BUILD_DIR
  cd $SOURCE_BUILD_DIR
  cmake .
  make VERBOSE=1
fi


COMMON_OPTIONS="-log_view \
-ksp_view \
-ksp_type cg \
-ksp_rtol 1.0e-8 \
-options_left"


ELASTICITY_OPTIONS="--problem_type elasticity \
-pc_type gamg \
-pc_gamg_coarse_eq_limit 1000 \
-mg_levels_ksp_type chebyshev \
-mg_levels_pc_type jacobi \
-mg_levels_esteig_ksp_type cg \
-matptap_via scalable"

POISSON_OPTIONS="--problem_type poisson"

# hypre is not (yet) available for complex petsc
if ! ( objdump -p $EXE | grep -q petsc_complex ); then
  POISSON_OPTIONS_REAL="-pc_type hypre \
  -pc_hypre_type boomeramg \
  -pc_hypre_boomeramg_strong_threshold 0.5"
fi

echo "==== elasticity weak scaling ($N_CPU processes with ndofs=$N_DOF_WEAK) ===="
mpirun -n $N_CPU $EXE --scaling_type weak --ndofs $N_DOF_WEAK $ELASTICITY_OPTIONS $COMMON_OPTIONS

echo "==== elasticity strong scaling ($N_CPU processes with ndofs=$N_DOF_STRONG) ===="
mpirun -n $N_CPU $EXE --scaling_type strong --ndofs $N_DOF_STRONG $ELASTICITY_OPTIONS $COMMON_OPTIONS

echo "==== poisson weak scaling ($N_CPU processes with ndofs=$N_DOF_WEAK) ===="
mpirun -n $N_CPU $EXE --scaling_type weak --ndofs $N_DOF_WEAK $POISSON_OPTIONS $POISSON_OPTIONS_REAL $COMMON_OPTIONS

echo "==== poisson strong scaling ($N_CPU processes with ndofs=$N_DOF_STRONG) ===="
mpirun -n $N_CPU $EXE --scaling_type strong --ndofs $N_DOF_STRONG $POISSON_OPTIONS $POISSON_OPTIONS_REAL $COMMON_OPTIONS
