#! /usr/bin/perl -w

use strict;
use Cwd;

my $mpiexec="mpiexec_mpt";
my $serpoe=undef;
my $mpmd=undef;
my $spmd=undef;
my $cmdfile=$ENV{MP_CMDFILE};
my $procs=undef;
my $cwd=getcwd();
my @pre=(); # ("$ENV{USHhwrf}/hwrf_clean_start.sh","$cwd");

sub set_pgmmode {
    my $pgmmode=shift @_;

    if(defined($pgmmode)) {
        if(lc($pgmmode) eq 'mpmd') {
            $mpmd=1; $spmd=undef;
            return;
        } elsif(lc($ENV{MP_PGMMODE}) eq 'spmd') {
            $mpmd=undef; $spmd=1;
            return;
        } else {
            return "Invalid PGMMODE: must be SPMD or MPMD, case-insensitive."
        }
    } else {
        $mpmd=undef; $spmd=1;
        return;
    }
}

sub set_proc_count {
    my $count=shift @_;
    if(defined($count)) {
        if($count=~/\A\s*0*([1-9][0-9]*)\s*\z/) {
            $procs=0+$1;
        } else {
            die "Invalid processor count \"$count\": must be an integer greater than 0.  POE IS ABORTING DUE TO INVALID PROCESSOR COUNT";
        }
    } else {
        die "No processor count available.  Please set \$TOTAL_TASKS or \$MP_PROCS or use --procs or -procs command-line options.  POE IS ABORTING BECAUSE THE PROCESSOR COUNT IS UNKNOWN";
    }
}

sub from_path {
    my $exe=shift @_;
    if($exe=~/^\// || $exe=~/^~/) {
        # Return absolute paths verbatim
        return $exe;
    } elsif($exe =~ /\//) {
        # Prepend the CWD to relative paths
        return $cwd."/".$exe;
    } else {
        # This exe has no path information
        if(-x "$exe") {
            # It exists in the current working directory.
            return $cwd."/".$exe;
        } else {
            # It is not in the CWD, so we must search the $PATH
            my $dir;
            foreach $dir(split(/:/,$ENV{PATH})) {
                if(-x "$dir/$exe") {
                    return $dir."/".$exe;
                }
            }

            die "$exe: does not exist in the current working directory ($cwd) and cannot find it in the path:\n$ENV{PATH}\nPOE IS ABORTING BECAUSE IT CANNOT FIND AN EXECUTABLE";
        }
    }
}

if(defined($ENV{USE_SERPOE})) {
    if(lc($ENV{USE_SERPOE}) eq 'yes') {
        $serpoe=1;
    } elsif(lc($ENV{USE_SERPOE}) eq 'no') {
        $serpoe=undef;
    } else {
        die "Invalid value for \$USE_SERPOE: must be \"yes\" or \"no\" (case-insensitive).  POE IS ABORTING DUE TO INVALID ARGUMENTS";
    }
}

set_pgmmode($ENV{MP_PGMMODEL});

if($ENV{MP_PROCS}) {
    set_proc_count($ENV{MP_PROCS});
} elsif($ENV{TOTAL_TASKS}) {
    set_proc_count($ENV{TOTAL_TASKS});
}

while($#ARGV>=0) {
    if($ARGV[0] =~ /\A-/) {
        if($ARGV[0] =~ /\A-?-pgmmodel\z$/) {
            die "$ARGV[0] requires an argument" unless $#ARGV>0;
            set_pgmmode($ARGV[1]);
            shift; shift;
        } elsif($ARGV[0] =~ /\A-?-cmdfile\z/) {
            die "$ARGV[0] requires an argument" unless $#ARGV>0;
            $cmdfile=$ARGV[1];
            shift; shift;
        } elsif($ARGV[0] =~ /\A-?-procs\z/) {
            die "$ARGV[0] requires an argument" unless $#ARGV>0;
            set_proc_count($ARGV[1]);
            shift; shift;
        } elsif($ARGV[0] eq '--') {
            # -- signifies the end of the argument list
            shift;
            last;
        } else {
            die "Unrecognized command line argument \"$ARGV[0]\".  POE IS ABORTING DUE TO AN INVALID COMMAND-LINE ARGUMENT";
        }
    } else {
        last;
    }
}

if(!defined($procs) || ! ($procs>=1)) {
    die "Unable to determine the number of MPI processes.  Please set \$TOTAL_TASKS or \$MP_PROCS to 1 or greater or use the -procs or --procs options.  POE IS ABORTING BECAUSE IT DOES NOT KNOW HOW MANY PROCESSES TO START";
}

my @argout=();

if($mpmd) {
    if(!defined($cmdfile) or $cmdfile=~/\A\s*\z/) {
        die "In MPMD mode, you must specify a command file via \$MP_CMDFILE, -cmdfile or --cmdfile.  POE IS ABORTING DUE TO MISSING COMMAND FILE";
    }
    if($#ARGV>=0) {
        die "In MPMD mode, you must not specify any command in the argument list.  POE IS ABORTING DUE TO UNEXPECTED EXTRA COMMAND-LINE ARGUMENTS";
    }

    if($serpoe) {
        @argout=('-n',$procs,@pre,from_path('serpoe'));
    } else {
        my($rank,$line,$prev);
        $prev='';

        open(CMD,"< $cmdfile") or die "$cmdfile: unable to open for reading: $!.  POE IS ABORTING BECAUSE IT CANNOT OPEN THE COMMAND FILE";

        my @addme=();
        for($rank=0;$rank<$procs;$rank++) {
            # Find the next non-blank line:
            while(1==1) {
                $line=<CMD>;
                chomp $line;
                if(!defined($line)) {
                    die "$cmdfile: cannot read enough lines for $procs ranks (got $rank lines): $!.  POE IS ABORTING BECAUSE IT CANNOT GET ENOUGH LINES FROM THE COMMAND FILE FOR ALL MPI RANKS";
                }
                next if $line=~/^\s*$/;    # ignore blank lines
                next if $line=~/^\s*#.*$/; # ignore comment lines
                $line=~s/^\s+//g;
                $line=~s/\s+$//g;
                last;
            }
            if($line eq $prev && $rank>0) {
                $addme[1]++;
            } else {
                if($#addme>=0) {
                    $prev=$line;
                    if($#argout>=0) {
                        push @argout,":",@addme;
                    } else {
                        push @argout,@addme;
                    }
                }
                $line=~/^\s*(\S+)(.*)/;
                if(-z $2) {
                    @addme=( '-n', 1, @pre, $1 );
                } else {
                    @addme=( '-n', 1, @pre, $1,$2 );
                }
                $addme[2+1+$#pre]=from_path($addme[2+1+$#pre]);
            }
        }
        if($#argout>=0) {
            push @argout,":",@addme;
        } else {
            push @argout,@addme;
        }
    }
} else {
    if($serpoe) {
        @argout=('-n',$procs,@pre,from_path('serpoe'),@ARGV);
    } else {
        @argout=('-n',$procs,@pre,@ARGV);
        $argout[2+1+$#pre]=from_path($argout[2+1+$#pre]);
    }
}

print STDERR "POE WILL NOW RUN: \"$mpiexec\", \"".join('","',@argout)."\"\n";

exec($mpiexec,@argout);
die "$mpiexec: UNABLE TO LAUNCH THE MPIEXEC COMMAND: $!.  POE IS ABORTING";