#!/bin/bash
#---------------------------------*- sh -*-------------------------------------
# ==  == ====== ====   ====    |
#                   \\     ||  | Multiphase Code Repository by HZDR
# ======   //   ||  || ===//   | Website: https://doi.org/10.14278/rodare.767
# ||  ||  //    ||  // || \\   | License: GPL-3.0-or-later
# ==  == ====== ====   ==  ==  |
#------------------------------------------------------------------------------
# License
#     This file is part of the Multiphase Code Repository by HZDR.
#
#     Copyright (C) 2025 by Helmholtz-Zentrum Dresden-Rossendorf e.V. (HZDR),
#     Website: https://hzdr.de
#
#     Multiphase Code Repository by HZDR is based on the free software for
#     computational fluid dynamics (CFD) from the OpenFOAM Foundation.
#     Copyright (C) 2025 by OpenFOAM Foundation, Website: https://openfoam.org
#
#     If you are interested in which files are original OpenFOAM Foundation
#     files, which OpenFOAM Foundation files were modified, and which files were
#     newly created, see FILES.md.
#
#     Multiphase Code Repository by HZDR is free software: you can redistribute
#     it and/or modify it under the terms of the GNU General Public License as
#     published by the Free Software Foundation, either version 3 of the
#     License, or (at your option) any later version.
#
#     Multiphase Code Repository by HZDR is distributed in the hope that it will
#     be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
#     Public License for more details.
#
#     You should have received a copy of the GNU General Public License along
#     with the Multiphase Code Repository by HZDR. If not, see
#     <http://www.gnu.org/licenses/>.
#
# Description
#     Alternative runParallel RunFunction that enables the use of Apptainer
#     containers in distributed parallel jobs using the Slurm workload manager.
#
#------------------------------------------------------------------------------

runCommand()
{
    wrapper="mpyshrun"
    if command -v "$wrapper" &> /dev/null; then
        $wrapper "$@"
    else
        echo "Please install multiphasepy:"
        echo "  pipx install multiphasepy --index-url \
https://codebase.helmholtz.cloud/api/v4/projects/8786/packages/pypi/simple"
    fi
}

runParallel()
{
    APP_RUN=
    LOG_IGNORE=false
    LOG_APPEND=false
    LOG_SUFFIX=
    nProcs=$(getNumberOfProcessors)

    # Parse options and executable
    while [ $# -gt 0 ] && [ -z "${APP_RUN}" ]; do
        key="$1"
        case "$key" in
            -append|-a)
                LOG_IGNORE=true
                LOG_APPEND=true
                ;;
            -overwrite|-o)
                LOG_IGNORE=true
                ;;
            -suffix|-s)
                LOG_SUFFIX=".$2"
                shift
                ;;
            -np|-n)
                nProcs="$2"
                shift
                ;;
            *)
                APP_RUN="$key"
                APP_NAME="${key##*/}"
                LOG_SUFFIX="${APP_NAME}${LOG_SUFFIX}"
                ;;
        esac

        shift
    done

    if [ -z "${APP_RUN}" ]
    then
        echo "Error: application not specified"
        return 1
    fi

    if [ -f log."${LOG_SUFFIX}" ] && [ "$LOG_IGNORE" = "false" ]
    then
        echo "${APP_NAME} already run on ${PWD}:" \
             "remove log file log.${LOG_SUFFIX} to re-run"
    else
        echo "Running ${APP_RUN} in parallel on ${PWD} using ${nProcs} \
            processes"

        if [ -n "${SINGULARITY_NAME:-}" ] \
        && [ -n "${SLURM_JOB_ID:-}" ] \
        && [ "${SLURM_NNODES:-}" -gt 1 ]; then
            if [ -n "${APPTAINER_NAME:-}" ]; then
                SHARENS="--sharens"
            else
                SHARENS=""
            fi

            if [ -n "${SINGULARITY_BIND:-}" ]; then
                BIND_COMMAND="-B ${SINGULARITY_BIND}"
            else
                BIND_COMMAND=""
            fi

            APP_OPTIONS=("$@")

            COMMAND="srun --mpi=pmi2 \
                        --jobid=${SLURM_JOB_ID} \
                        --ntasks=${SLURM_NTASKS} \
                        singularity exec ${SHARENS} ${BIND_COMMAND} \
                        ${SINGULARITY_CONTAINER} ${APP_RUN} -parallel"

            if [ "${LOG_APPEND}" = "true" ]; then
                # Expansion of here document happens on the client side on
                # purpose https://www.shellcheck.net/wiki/SC2087
                # shellcheck disable=SC2087
                ssh "${SLURM_SUBMIT_HOST}" <<EOF > /dev/null 2>&1
                    $(declare -p APP_OPTIONS)
                    cd ${PWD}
                    ${COMMAND} "\${APP_OPTIONS[@]}" \
                        < /dev/null >> "log.${LOG_SUFFIX}" 2>&1
EOF
            else
                # shellcheck disable=SC2087
                ssh "${SLURM_SUBMIT_HOST}" <<EOF > /dev/null 2>&1
                    $(declare -p APP_OPTIONS)
                    cd ${PWD}
                    ${COMMAND} "\${APP_OPTIONS[@]}" \
                        < /dev/null > "log.${LOG_SUFFIX}" 2>&1
EOF
            fi
        else
            if [ "${LOG_APPEND}" = "true" ]; then
                ( mpirun -np "${nProcs}" "${APP_RUN}" -parallel "$@" \
                    < /dev/null >> "log.${LOG_SUFFIX}" 2>&1 )
            else
                ( mpirun -np "${nProcs}" "${APP_RUN}" -parallel "$@" \
                    < /dev/null > "log.${LOG_SUFFIX}" 2>&1 )
            fi
        fi
    fi
}

#------------------------------------------------------------------------------
