#! /bin/env bash

# This script is a poe-lookalike script that will allow the HWRF scripts
# to work on Jet without having to replace all the "poe" calls with
# mpirun, mpiexec or whatever else.

# This is written for Jet, but should be easy to extend to other
# clusters by changing the mpiexec calls near the bottom of this
# script to something else (mpirun_rsh, mpirun, etc.)

function message {
    echo "$*" 1>&2
}
function die {
    set +x
    message "$*"
    exit 2
}
function usage {
    cat<<EOF 1>&2
This script pretends to be the IBM poe command.  However, it is
considerably stupider than poe, and only understands a few arguments.
Also, you MUST set \$TOTAL_TASKS or \$MP_PROCS to the total number of
processors to use, or specify the -procs argument.

Call it like this:

    export MP_PGMMODEL=mpmd  # case-insensitive
    export MP_CMDFILE=/path/to/command/file
    poe

 OR:

    export MP_PGMMODEL=spmd  # set to anything but mpmd (case-insensitive)
    poe /path/to/my/program (arguments ...)

 UNDERSTOOD OPTIONS:

    -procs number of processors (same as \$MP_PROCS)
    -pgmmodel MPMD or SPMD (same as \$MP_PGMMODEL)
    -cmdfile  /path/to/command/file
         (same as \$MP_CMDFILE, implies -pgmmodel MPMD)
EOF
    if [[ "$#" -gt 0 ]] ; then
        echo
        for x in "$@" ; do message "$x" ; done
        exit 2
    fi
    exit 0
}

function arg {
    if [[ "$#" -lt 2 ]] ; then
        usage "SCRIPT IS ABORTING: $1 REQUIRES AN ARGUMENT"
    fi
}


argerr=NO
go=YES
while [[ "$#" -gt 0  && "$go" == YES ]] ; do
    case "$1" in
        -pgmmodel|--pgmmodel)
            export MP_PGMMODEL="$2"
            arg "$@" ; shift 2 ;;
        -cmdfile|--cmdfile)
            export MP_PGMMODEL=mpmd
            export MP_CMDFILE="$2"
            arg "$@" ; shift 2 ;;
        -procs|--procs)
            if [[ ! ( "$2" -gt 0 ) ]] ; then
                die "UH-OH!! You used -procs \"$2\" but \"$2\" is not a number greater than zero.  I'm giving up now."
            fi
            export MP_PROCS="$2"
            arg "$@" ; shift 2 ;;
        --) shift ; go=NO ;;
        -*) argerr="SCRIPT IS ABORTING: UNRECOGNIZED OPTION $1" ; go=NO ;;
        *) go=NO ;;
    esac
done
if [[ "$argerr" != NO ]] ; then
    usage "$argerr"
fi

if [[ "$MP_PROCS" != "$TOTAL_TASKS" && ! -z "$TOTAL_TASKS" ]] ; then
    message "Overriding number of processes \$TOTAL_TASKS with \$MP_PROCS=$MP_PROCS."
fi

set -x

export TOTAL_TASKS="${MP_PROCS:-$TOTAL_TASKS}"

if [[ ! ( "$TOTAL_TASKS" -gt 0 ) ]] ; then
    die "UH-OH!! Your \$TOTAL_TASKS and \$MP_PROCS are both undefined and you did not specified -procs!! I don't know what to do now.  I guess I'll just exit.  Sorry for being so bad at this -- I really try to be as good as the real IBM poe, but I just can't keep up with the real poe's genius."
fi

case Q"$MP_PGMMODEL" in
    Q[mM][pP][mM][dD])
        if [[ -z "$MP_CMDFILE" ]] ; then
            die "BY GOLLY!!  You have not set the \$MP_CMDFILE variable.  Either fill one in, or use SPMD.  ABORTING"
        fi

        if [[ ! -s "$MP_CMDFILE" ]] ; then
            die "GOLLY GEE!!  Your \$MP_CMDFILE (\"$MP_CMDFILE\") is empty or non-existant.  Don't know what to do.  ABORTING"
        fi

        if [[ ! -z "$TOTAL_TASKS" ]] ; then
            them_lines=$( cat "$MP_CMDFILE" | perl -ne '/^\s*$/ and next; $x++; END { print "$x\n" }' )
            if [[ "$them_lines" != "$TOTAL_TASKS" ]] ; then
                die "WHOA NELLY!!  Your \$MP_CMDFILE (\"$MP_CMDFILE\") has $them_lines lines, but \$TOTAL_TASKS=$TOTAL_TASKS!!  You had better fix that.  ABORTING"
            fi
        fi

        # Run sync to ensure we're using the most up-to-date copy of
        # the command file (to get around Jet's aggressive file
        # caching):
        sync

        # For MPMD, we have to do some extra work.  The below perl
        # command generates an "mpiexec ..." command that the 'exec'
        # command will execute.
        exec $( perl -ne '
            BEGIN {
                $nj=0;
            };
            chomp;              # remove end-of-line char
            /^\s*$/ and next;   # discard blank lines
            if($nj<$ENV{TOTAL_TASKS}) {
                $nj++;
                if($#c>=0 && $c[$#c] eq $_) {
                    # This command is the same as the previous, so
                    # just increment that commands -np N instead of
                    # inserting a new command:
                    $n[$#n]++;
                } else {
                    # Add this command to the list with -np 1
                    push @c, $_;
                    push @n, 1;
                }
            }
            END {
                # Now that we have read in the input, generate a command.
                print "mpiexec ";
                for($i=0; $i<=$#n; $i++) { # loop over all 
                    print "-np $n[$i] -envall $c[$i]";
                    if($i<$#n) {
                        print " : ";
                    }
                }
            }
        ' < "$MP_CMDFILE" )
        ;;
    Q[sS][pP][mM][dD]|Q)
        # We're doing SPMD or no value specified (which implies SPMD).
        if ( echo "$1" | grep / > /dev/null ) ; then
            # Caller specified the path to an executable, so make sure
            # that executable exists, is non-empty and is executable
            if [[ ! -e "$1" ]] ; then
                message "SLIVERING SNAKES!! Your program \"$1\" does not seem to exist!  I'll try executing it anyway but this will probably fail..."
            elif [[ ! -s "$1" ]] ; then
                message "WHOPPING WILLABIES!! Your program \"$1\" seems to be empty!  I'll try executing it anyway but this will probably fail..."
            elif [[ ! -x "$1" ]] ; then
                message "BUCKING BRONCOS!! Your program \"$1\" does not seem to be executable.  I'll try executing it anyway but this will probably fail..."
            fi
        fi
        if [[ "$TOTAL_TASKS" -gt 1 ]] ; then
            exec mpiexec -envall -np "$TOTAL_TASKS" $*
        elif [[ "$TOTAL_TASKS" == 1 ]] ; then
            message "WARNING: RUNNING ONLY ONE TASK SO RUNNING IT DIRECTLY INSTEAD OF USING MPIRUN (TO WORK AROUND A JET BUG)."
            exec $*
            die "OOPES!! Unable to run $*.  ABORTING"
        else
            message "WARNING: \$TOTAL_TASKS IS NOT SPECIFIED OR IS NOT A POSITIVE NUMBER, SO I WILL NOT SPECIFY -np WHEN RUNNING mpiexec"
            exec mpiexec -envall $*
        fi
esac

# Should never get here.
exit 99