#! /bin/ksh

set -x

# This script hides the details of launching an MPI program, or
# multiple serial programs in parallel.  It provides a cross-platform
# implementation that works on Cirrus, Stratus, Vapor, Jet and Zeus
# (and hopefully, one day, GAEA).  Currently it is dependent on a
# partial poe implementation, and one is provided on Jet and Zeus in
# other directories not given to NCO.

# Default values
mpiwrap=no # yes=run serial programs, wrapped in MPI
runshell=no # yes=run shell scripts, wrapped in MPI; implies mpiwrap=yes
mode=spmd # spmd (all ranks have same program) or mpmd (multiple programs)
cmdfile=cmdf # command file for MPMD mode (ignored in SPMD mode)

# Default values for task geometry (machine list modification).
# Ignored when using LoadLeveler and POE due to limitations of the
# IBM MPI implementation

endtasks='' # number of tasks per machine at end of list (3:8:4 means
            # 4 on last machine, 8 on second to last, 3 on third to last)
starttasks='' # Similar, but for start of list (2:8:1 means 2 on first
              # machine, 8 on second, 1 on third)
midtasks='' # number of tasks to distribute inside the remaining machines
grid='' # specify instead of "midtasks," this distributes the tasks as
        # a grid: 3x4:10x10 makes a 10x10 grid of machines, with a 3x4
        # grid of processors within.
tweak_hostfile=no # =yes if any of endtasks, starttasks, midtasks or
                  # grid are specified

function fail {
    # This function terminates the job with an error message:
    sh ${utilscript}/setup.sh
    err=911
    err_exit "$*"
}

function needarg {
    # This function is used during argument parsing to check and see
    # if there are at least two arguments left to parse on the command
    # line.  The first argument remaining to be parsed is assumed to
    # be a dash option (-n, -mpmd, etc.) in $opt, and the second is
    # the argument to that option.  The needarg function is then
    # called with $# as its only argument, and if $# ($1 in this
    # function) is less than 2, then an argument was not given to the
    # dash option, and we must terminate with an error message.
    if [[ "$1" -lt 2 ]] ; then
        fail "Option $opt to hwrf_run.sh requires an argument."
    fi
}

# Parse arguments until the first argument dealing with the command or
# commands to launch:
go=yes
while [[ $go == yes && "$#" -gt 0 ]] ; do
    opt="$1"
    case "$opt" in
        --)                             shift 1 ; go=no ;;
        -mpiwrap)     mpiwrap=yes     ; shift 1         ;;
        -runshell)    runshell=yes
                      mpiwrap=yes     ; shift 1         ;;
        -mpmd)        mode=mpmd       ; shift 1         ;;
        -spmd)        mode=spmd       ; shift 1         ;;
        -cmdfile)     needarg "$#"
                      cmdfile="$2"    ; shift 2         ;;
        -endtasks)    needarg "$#"
                      tweak_hostfile=yes
                      endtasks="$2"   ; shift 2         ;;
        -starttasks)  needarg "$#"
                      tweak_hostfile=yes
                      starttasks="$2" ; shift 2         ;;
        -midtasks)    needarg "$#"
                      tweak_hostfile=yes
                      midtasks="$2"   ; shift 2         ;;
        -grid)        needarg "$#"
                      tweak_hostfile=yes
                      grid="$2"       ; shift 2         ;;
        -n)                                       go=no ;;
        -*)
            fail "ABORTING: UNRECOGNIZED ARGUMENT $opt" ;;
        *) go=no
    esac
done



# If requested, modify the host file to apply task geometry or fewer
# processors per machine on the first few or last few machines for I/O
# servers, high-memory coupling tasks, or tasks that cause a lot of
# system background activity.  This is impossible on CCS due to
# limitations of LoadLeveler and the IBM MPI implementation, but it is
# used elsewhere with great benefit to runtimes:

if [[ "$tweak_hostfile" == yes && "$PARAFLAG" == YES ]] ; then
    set +x
    if [[ ! -z "$midtasks" && ! -z "$grid" ]] ; then
        fail "ABORTING: TRIED TO TWEAK HOSTFILE BUT BOTH -grid AND -midtasks UNSPECIFIED"
    fi
    tweak_hostfile -d ${endtasks:+-e $endtasks} ${starttasks:+-s $starttasks} \
        ${grid:+-g $grid} ${midtasks:+-p $midtasks}
    set -x
fi



# Generate the MPMD command file if needed, and set the relevant
# environment variables:

export MKL_NUM_THREADS=1
export OMP_NUM_THREADS=1
export USE_SERPOE="$mpiwrap"
export SERPOE_SHELL="$runshell"

if [[ "$mode" == mpmd ]] ; then
    export MP_PGMMODEL=mpmd
    export MP_CMDFILE=${cmdfile:-cmdf}

    if [[ "$#" -gt 0 ]] ; then
        # We are being asked to generate the command file.
        rm -f "$MP_CMDFILE"
        
        # Generate the command file from an argument list like this:
        #  -n 1 coupler : 
        count=1
        command=""
        while [[ "$#" -gt 0 ]] ; do
            opt=$1
            case "$opt" in
                -n) needarg "$#"
                    count="$2"
                    shift 2
                    ;;
                :)
                    # end of this command
                    for x in $( seq 1 $count ) ; do
                        echo "$command" >> "$MP_CMDFILE"
                    done
                    count=1
                    command=''
                    shift 1
                    ;;
                *)  command="${command:+$command }$1"
                    shift 1
            esac
        done
        if [[ ! -z "$command" ]] ; then
            # append the last command to the command file
            for x in $( seq 1 $count ) ; do
                echo "$command" >> "$MP_CMDFILE"
            done
        fi
    fi
else
    export MP_PGMMODEL=spmd
    unset MP_CMDFILE
fi



# Start poe or the poe-lookalike
if [[ "$mode" == mpmd ]] ; then
    poe
else
    poe "$@"
fi