#!/bin/bash
#
# Conv-host for MPI:
#  Translates +pN-style conv-host options into
# mpirun -npN options.

export FROM_CHARMRUN='1'

args=()

# Arguments that need to be passed to the charm application, not mpirun (e.g. ++quiet)
charm_args=()
pes=0
nodes=0
ppn=1
QUIET=0
XTERM='xterm'
DEBUG=0
DEBUG_NO_PAUSE=0
DEBUGGER='gdb'

runCmd()
{
  [[ $QUIET -eq 0 ]] && echo "charmrun>" "$@"
  "$@"
}

while [[ $# -ne 0 ]]
do
	case "$1" in
	+n|++n|++np)
		nodes="$2"
		shift
		;;
	+n[0-9]*)
		nodes="${1#\+n}"
		;;
	++n[0-9]*)
		nodes="${1#\+\+n}"
		;;
	++np[0-9]*)
		nodes="${1#\+\+np}"
		;;
	+ppn|++ppn)
		args+=(+ppn "$2")
		ppn="$2"
		shift
		;;
	+ppn[0-9]*)
		args+=("$1")
		ppn="${1#\+ppn}"
		;;
	++ppn[0-9]*)
		args+=("$1")
		ppn="${1#\+\+ppn}"
		;;
	+p|++p)
		pes="$2"
		shift
		;;
	+p[0-9]*)
		pes="${1#\+p}"
		;;
	++p[0-9]*)
		pes="${1#\+\+p}"
		;;
	+pemap)
	    args+=("$1" "$2")
	    shift
	    ;;
	-machinefile)
	    args=("$1" "$2" "${args[@]}")
	    shift
	    ;;
	++quiet)
	    QUIET=1
	    ;;
	++no-quiet)
	    QUIET=0
	    ;;
	++local|++no-local)
		# Ignored (unsupported by MPI layer)
		;;
	++xterm)
		XTERM="$2"
		shift
		;;
	++debug)
		DEBUG=1
		;;
	++no-debug)
		DEBUG=0
		;;
	++debug-no-pause)
		DEBUG_NO_PAUSE=1
		;;
	++no-debug-no-pause)
		DEBUG_NO_PAUSE=0
		;;
	++debugger)
		DEBUGGER="$2"
		shift
		;;
	*)
	    args+=("$1")
	    ;;
    esac
    shift
done

[[ $QUIET -eq 1 ]] && charm_args+=(++quiet)

args+=("${charm_args[@]}")

if [[ "$DEBUG" = '1' || "$DEBUG_NO_PAUSE" = '1' ]]
then
  DEBUG_RUN=(-ex r)
  DEBUG_POSTFIX='--args'
  if [[ "$DEBUGGER" = 'lldb' ]]
  then
    DEBUG_RUN=(-o r)
    DEBUG_POSTFIX='--'
  fi

  args=("$DEBUG_POSTFIX" "${args[@]}")
  [[ "$DEBUG_NO_PAUSE" = '1' ]] && args=("${DEBUG_RUN[@]}" "${args[@]}")
  args=("$XTERM" -e "$DEBUGGER" "${args[@]}")
fi

if [[ $nodes = 0 ]]
then
  [[ $pes = 0 ]] && pes=1
  if [[ $(( pes % ppn )) -ne 0 ]]
  then
    printf "Charmrun> Error: ++ppn %s (number of PEs per node) does not divide +p %s (number of PEs).\n" "$ppn" "$pes"
    exit 1
  else
    nodes=$(( pes / ppn ))
  fi
elif [[ $pes != 0 && $pes != $(( nodes * ppn )) ]]
then
  printf "Charmrun> Error: +n/++np %s times ++ppn %s does not equal +p %s.\n" "$nodes" "$ppn" "$pes"
  exit 1
fi

[[ $QUIET -eq 0 ]] && echo && echo "Running as $nodes OS processes: ${args[*]}"


mpiexec_cmd="$(command -v mpiexec 2>/dev/null)"
if [[ -n "$mpiexec_cmd" ]]
then
    runCmd "$mpiexec_cmd" -n "$nodes" "${args[@]}"
else
    echo "Don't know how to run MPI program."
    exit 1
fi
