#!/bin/bash
#
# Copyright (c) 2014      Artem Polyakov <artpol84@gmail.com>
# $COPYRIGHT$
#
# Additional copyrights may follow
#
# $HEADER$

function get_jobid()
{
    if [ -n "$SLURM_JOBID" ]; then
        echo "$SLURM_JOBID"
        return
    elif [ -n "$PBS_JOBID" ]; then
        echo "$PBS_JOBID"
        return
    else
        echo "$$"
        return
    fi
}

function extract_nodeset_opts()
{
    short_opts="H:h:"
    long_opts="host:"
    opts=""

    while [ -n "$1" ]; do
        case "$1" in
            -H | -host | --host )
                opts=$opts" $1 $2";
                shift 2;;
            -hostfile | --hostfile )
                opts=$opts" $1 $2";
                shift 2;;
            -machinefile | --machinefile )
                opts=$opts" $1 $2";
                shift 2;;
            *)
                shift 1;;
        esac
    done

    echo "$opts"
}


ompi_instdir=`dirname $0`
jobid=`get_jobid`
syncfile=`pwd`"/ompi_clock_sync_data.$jobid"
tmp_timings=`pwd`"/ompi_timing_temp_file.$jobid"
tmp_out=`pwd`"/ompi_mpirun_prof.$jobid"

timing_bkp=$OMPI_MCA_opal_timing_output
export OMPI_MCA_opal_timing_output=$tmp_timings

opts=`extract_nodeset_opts $@`
${ompi_instdir}/mpirun --npernode 1 $opts ${ompi_instdir}/mpisync -o $syncfile >$tmp_out 2>&1

export OMPI_MCA_opal_timing_output=$timing_bkp
export OMPI_MCA_opal_timing_sync_file=$syncfile

# Remove old output
rm -f $OMPI_MCA_opal_timing_output

# Run a program of interest
${ompi_instdir}/mpirun $@

if [ -n "$timing_bkp" ]; then
    ${ompi_instdir}/ompi_timing_post $timing_bkp $timing_bkp.post
fi

# Cleanup
rm -f $syncfile
rm -f $tmp_timings
rm -f $tmp_out