# ICON
#
# ---------------------------------------------------------------
# Copyright (C) 2004-2026, DWD, MPI-M, DKRZ, KIT, ETH, MeteoSwiss
# Contact information: icon-model.org
# See AUTHORS.TXT for a list of authors
# See LICENSES/ for license information
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------

###
# START of run/add_START_MODEL_function
# This file is included via run/create_target_header on the DWD NEC machine

START_MODEL_function()
{ # START_MODEL is a function. This has the advantage that the variables within
  # are just replaced at execution time, i.e. they can be modified by the code
  # before.
  #
  # Environment variables
  # ---------------------
  ## required:
  # PPN: total number of MPI procs. (on vector engines) per job
  # NE: number of vector engines
  #
  ## optional to refine VE/VH configuration:
  # proc0_shift: if > 0 allocate proc0 on VH (for all components), else allocate
  #     all procs on VE. (default: 0)
  # ICON_COMPONENT${c}_proc0_shift: if == 0, don't allocate proc0 on VH for
  #     component c. (This setting is rather a workaround that is necessary
  #     until all components support proc0_shit.) (default: $proc0_shift)
  # ICON_COMPONENT${c}_VE_procs: Number of vector engine processes for ICON
  #     component c, (default: PPN - (sum of VE_procs of all other components)
  #     for component 1, 0 for all others)
  # ICON_COMPONENT${c}_VH_procs: Number of vector host processes for ICON
  #     component 1, (default: 0)
  #     This excludes proc0. Component 1 always allocates a VH proc for proc0.
  # ICON_COMPONENT${c}_VH_bin: working-dir relative path to the VH binary of
  #     component c, (default: $MODEL_SCAL)
  # ICON_COMPONENT${c}_VE_bin: working-dir relative path to the VE binary of
  #     component c, (default: $MODEL)
  # venum_lhost: Number of VE per logical host (default: 2)

    # Check mandatory variables
    : ${PPN:?START_MODEL_function: PPN is not set}
    : ${NE:?START_MODEL_function: NE is not set}
    : ${MODEL_SCAL:?START_MODEL_function: MODEL_SCAL is not set}
    : ${MODEL:?START_MODEL_function: MODEL is not set}
    : ${START:?START_MODEL_function: START is not set}

    # Note from `man bash`: `:` is a shell builtin command and has no effect.
    # "the command does nothing beyond expanding arguments and performing any
    # specified redirections."

    # Set optional variables
    local proc0_shift=${proc0_shift:-0}
    local venum_lhost=${venum_lhost:-2}

    declare -i c=1 nveprocs_excl_first=0 any_nonstandard=0
    declare -a proc0_shift_c VH_procs_c VE_procs_c VH_bin_c VE_bin_c

    # Get the number of mentioned components and their allocations.
    while true; do
        declare -n n_proc0_shift="ICON_COMPONENT${c}_proc0_shift"
        declare -n n_VH_procs="ICON_COMPONENT${c}_VH_procs"
        declare -n n_VE_procs="ICON_COMPONENT${c}_VE_procs"
        declare -n n_VH_bin="ICON_COMPONENT${c}_VH_bin"
        declare -n n_VE_bin="ICON_COMPONENT${c}_VE_bin"

        # Check if component $c is mentioned at all (component 1 always is, so it can implicitly
        # take the remaining procs)
        if (( c > 1)) && [[ -z ${n_proc0_shift+x} && -z ${n_VH_procs+x} && -z ${n_VE_procs+x} ]]; then
          break
        fi

        proc0_shift_c+=( "${n_proc0_shift:-$proc0_shift}" )
        VH_procs_c+=( "${n_VH_procs:-0}" )
        any_nonstandard=$(( any_nonstandard || proc0_shift_c[c-1] > 0 || VH_procs_c[c-1] > 0 ))

        if (( c == 1 )); then
            VE_procs_c+=( "${n_VE_procs:--1}" )
        else
            VE_procs_c+=( "${n_VE_procs:-0}" )
            nveprocs_excl_first=$(( nveprocs_excl_first + VE_procs_c[c-1] ))
        fi

        VH_bin_c+=( "${n_VH_bin:-$MODEL_SCAL}" )
        VE_bin_c+=( "${n_VE_bin:-$MODEL}" )

        if [[ -n ${n_VE_bin+x} ]]; then
            # A different VE binary was specified. Switch to full allocation specification.
            any_nonstandard=1
        fi

        c=$(( c + 1 ))
    done

    # Number of components
    declare -ri NUM_COMPONENTS=$(( c - 1 ))

    # First component takes remaining procs if nothing was specified.
    if (( VE_procs_c[0] < 0 )); then
        VE_procs_c[0]=$(( PPN - nveprocs_excl_first ))
    fi

    local procs_per_ve=$(( (PPN + NE - 1)/NE )) # this should be about the same as $CPE
    local number_vector_hosts=$(( (NE + venum_lhost - 1)/venum_lhost))

    if (( any_nonstandard )); then
        local start_options=""

        local split_ve=-1
        local remaining_cores=0

        for (( c=0; c < NUM_COMPONENTS; c+=1 )); do
            local _proc0_shift="${proc0_shift_c[c]}"
            local _VH_procs="${VH_procs_c[c]}"
            local _VE_procs="${VE_procs_c[c]}"
            local _VH_bin="${VH_bin_c[c]}"
            local _VE_bin="${VE_bin_c[c]}"
            local cc=$(( c + 1 ))

            # Compute PPN for even distribution of IO and restart procs
            # (number_vector_hosts is the number of allocated VHs)
            local VH_ppn=$(( (_VH_procs + number_vector_hosts - 1)/number_vector_hosts ))

            if (( _proc0_shift > 0 )); then
                # distribute proc0s across hosts if possible
                local proc0_vh=$(( c % number_vector_hosts ))
                echo "Component $cc: Allocate proc0 on VH $proc0_vh ($proc0_shift proc(s))"
                start_options="$start_options : -vh -node $proc0_vh -np $proc0_shift -env OMP_NUM_THREADS 1 ${_VH_bin}"
            fi

            local ve_procs_left=$_VE_procs

            if (( remaining_cores > 0 && ve_procs_left > 0 )); then
                # distribute leftover cores from previous component
                local split_ve_procs=$remaining_cores
                if (( ve_procs_left < remaining_cores )); then
                    split_ve_procs=$ve_procs_left
                fi
                ve_procs_left=$(( ve_procs_left - split_ve_procs ))
                remaining_cores=$(( remaining_cores - split_ve_procs ))

                echo "Component $cc: Allocate $split_ve_procs procs on the split VE $split_ve"
                start_options="$start_options : -venode -node $split_ve -nnp $split_ve_procs -np $split_ve_procs -env OMP_NUM_THREADS ${OMP_NUM_THREADS:-1} ${_VE_bin}"
            fi

            local ves=$(( (ve_procs_left + procs_per_ve - 1)/procs_per_ve ))

            if (( ves > 0 )); then
                local start_ve=$(( split_ve + 1 ))
                local end_ve=$(( start_ve + ves - 1 ))

                echo "Component $cc: Allocate $ve_procs_left procs on VEs ${start_ve}-${end_ve} ($procs_per_ve per VE)"
                start_options="$start_options : -venode -node ${start_ve}-${end_ve} -nnp $procs_per_ve -np $ve_procs_left -env OMP_NUM_THREADS ${OMP_NUM_THREADS:-1} ${_VE_bin}"

                split_ve=$end_ve
                remaining_cores=$(( ves*procs_per_ve - ve_procs_left ))
            else
                : # No additional VEs needed. Keep split_ve and remaining_cores as is
            fi

            if (( _VH_procs > 0 )); then
                echo "Component $cc: Allocate $_VH_procs VH procs distributed over all hosts (up to $VH_ppn procs per host)"
                start_options="$start_options : -vh -nnp $VH_ppn -np $_VH_procs -env OMP_NUM_THREADS 1 ${_VH_bin}"
            fi
        done

        if (( remaining_cores > 0 )); then
            echo "$remaining_cores unassigned VE cores remain on VE $split_ve"
        fi

    else # Only VEs, default binary.
        local start_options="-x -venode -node 0-$((NE - 1)) -np $PPN -env OMP_NUM_THREADS ${OMP_NUM_THREADS:-1} $MODEL"
    fi
    $START -v ${start_options# :}
    return $?
}
START_MODEL=START_MODEL_function

#
# END of run/add_START_MODEL_function
###
