#!/bin/bash
IS_CRAY=no
IS_MIC=no
IS_STATIC=no

# First argument is set when this script is called from MAP or Performance
# Reports. This should be the hostname. Ignore argument if it begins with '-'.
if [[ -z "$1" || $1 = -* ]]; then
  MAP_WRAPPER_POSTFIX="-`hostname`"
else
  MAP_WRAPPER_POSTFIX="-$1"
fi

# Second argument is set when this script is called from MAP or Performance
# Reports. This should be the pid of the MAP or Performance Reports process.
# Ignore argument if it begins with '-'.
if [[ -z "$2" || $2 = -* ]]; then 
  MAP_WRAPPER_PID=""
else
  MAP_WRAPPER_POSTFIX="${MAP_WRAPPER_POSTFIX}-"
  MAP_WRAPPER_PID="$2"
fi

# Third argument is set when this script is called from MAP or Performance
# Reports. This should be the mpirun-like script that was called. 
# Ignore argument if it begins with '-'.

if [[ -z "$3" || $3 = -* ]]; then
  MPIEXECUTABLE="mpirun"
else
  MPIEXECUTABLE="$3"
fi

if [[ -z "$4" || $4 = -* ]]; then 
  MPINAME="None"
else
  MPINAME="$4"
fi

for var in "$@"
do
  if [ "$var" = "--cray" ]; then
    IS_CRAY=yes
  fi

  if [ "$var" = "--mic" ]; then
    IS_MIC=yes
  fi

  if [ "$var" = "--static" ]; then
    IS_STATIC=yes
  fi
done

# Check if FORGE_TOOLS_PATH is set. Don't set this in your environment - the
# MAP or Performance Reports launch script will have set it.
if [ -z "${FORGE_TOOLS_PATH}" ]; then
    # If FORGE_TOOLS_PATH has not been set, work out where MAP or Performance
    # Reports is located from path to this script
    FORGE_TOOLS_PATH=`dirname "$0"`/../..
fi

if [ "x$FORGE_CONFIG_DIR" = "x" ]; then
  # if empty, try FORGE_TOOLS_CONFIG_DIR
  FORGE_CONFIG_DIR="$FORGE_TOOLS_CONFIG_DIR"
fi
if [ "x$FORGE_CONFIG_DIR" = "x" ]; then
  # if empty, try ALLINEA_CONFIG_DIR
  FORGE_CONFIG_DIR="$ALLINEA_CONFIG_DIR"
fi
if [ "x$FORGE_CONFIG_DIR" = "x" ]; then
  # if empty, try ALLINEA_TOOLS_CONFIG_DIR
  FORGE_CONFIG_DIR="$ALLINEA_TOOLS_CONFIG_DIR"
fi
if [ "x$FORGE_CONFIG_DIR" = "x" ]; then
  # if still empty, set to default
  FORGE_CONFIG_DIR="$HOME/."allinea
fi

if [ ! -d "$FORGE_CONFIG_DIR" ]; then
  mkdir "$FORGE_CONFIG_DIR"
  if [ $? -ne 0 ]; then
    echo "Failed to create directory '$FORGE_CONFIG_DIR'"
    echo "To use a different directory, set the FORGE_CONFIG_DIR environment variable."
    exit $?
  fi
fi

MAP_WRAPPER_DIR=${FORGE_CONFIG_DIR}/wrapper
if [ ! -d "$MAP_WRAPPER_DIR" ]; then
  mkdir "$MAP_WRAPPER_DIR"
  if [ $? -ne 0 ]; then
    exit $?
  fi
fi

# Add -mic suffix when cross-compiling Xeon Phi libraries
if [ "${IS_MIC}" = "yes" -a -z "${MAP_WRAPPER_PID}" ]; then
  MAP_WRAPPER_POSTFIX="${MAP_WRAPPER_POSTFIX}-mic"
fi

MAP_WRAPPER_LIB=${MAP_WRAPPER_DIR}/libmap-sampler-pmpi${MAP_WRAPPER_POSTFIX}${MAP_WRAPPER_PID}.so
MAP_WRAPPER_CODE=${MAP_WRAPPER_DIR}/libmap-sampler-pmpi${MAP_WRAPPER_POSTFIX}${MAP_WRAPPER_PID}.c
MAP_WRAPPER_OBJECT=${MAP_WRAPPER_DIR}/libmap-sampler-pmpi${MAP_WRAPPER_POSTFIX}${MAP_WRAPPER_PID}.o
MAP_WRAPPER_STATIC_CODE_DIR=${MAP_WRAPPER_DIR}/libmap-sampler-pmpi${MAP_WRAPPER_POSTFIX}${MAP_WRAPPER_PID}

# Directory in which to add all needed precompiledwrappers
PRELOAD_WRAPPER_DIR=${FORGE_TOOLS_PATH}/map/wrapper/precompiled

# Remove the old sampler MPI wrapper library (if it exists). The MPI
# implementation in the PATH may have changed since the wrapper was last
# built.
echo "Deleting old wrapper"
rm -f ${MAP_WRAPPER_LIB}* ${MAP_WRAPPER_CODE}

# Copies everything that is required for the wrapper to work and renames it to how build_wrapper creates them
function renamecopy(){
  echo "Renaming precompiled wrappers and copying them"
  cp ${PRELOAD_WRAPPER_DIR}/$1/wrapper/libmap-sampler-pmpi-precompiled.so.1.0.0 $MAP_WRAPPER_LIB.1.0.0
  ln -s $(basename $MAP_WRAPPER_LIB.1.0.0) $MAP_WRAPPER_LIB.1.0
  ln -s $(basename $MAP_WRAPPER_LIB.1.0) $MAP_WRAPPER_LIB.1
  ln -s $(basename $MAP_WRAPPER_LIB.1) $MAP_WRAPPER_LIB

  echo "Renamed wrapper: $1"
  echo "MPI wrapper copied: ${MAP_WRAPPER_LIB}"
  exit 0
}

# Discovers the GNU version loaded in Cray
# Expected environment variables:
# $ echo $PE_MPICH_DEFAULT_GENCOMPS_GNU -> 51 49
# $ echo $GNU_VERSION                   -> 7.3.0
# $ echo $PE_MPICH_GENCOMPILERS_GNU     -> 9.1
function discover_cray_gnu_version(){
  local VERSION=""
  if [[ -n "$PE_MPICH_DEFAULT_GENCOMPS_GNU" ]]; then
    VERSION=${PE_MPICH_DEFAULT_GENCOMPS_GNU:0:2}
  elif [[ -n "$PE_MPICH_GENCOMPILERS_GNU" ]]; then
    VERSION=${PE_MPICH_GENCOMPILERS_GNU:0:1}${PE_MPICH_GENCOMPILERS_GNU:2:1}
  elif [[ -n "$GNU_VERSION" ]]; then
    if [[ ${GNU_VERSION:2:1} == "." ]]; then
      # GCC XX.X
      VERSION=${GNU_VERSION:0:2}${GNU_VERSION:3:1}
    else
      # GCC X.X
      VERSION=${GNU_VERSION:0:1}${GNU_VERSION:2:1}
    fi
  else
     echo "Can't find GCC for the MPI version"
     exit 1
  fi

  SUPPORTED_GCC=()
  if [[ "${CRAY_MPICH_MAJOR_VERSION}" -eq "7" ]] ; then
    SUPPORTED_GCC=(51 71 73)
  elif [[ "${CRAY_MPICH_MAJOR_VERSION}" -eq "8" ]] ; then
    SUPPORTED_GCC=(91)
  fi

  if [[ ! "${SUPPORTED_GCC[@]}" =~ "${VERSION}" ]]; then
    echo "Unsupported MPI version, GCC version $VERSION not supported"
    exit 1
  fi

  CRAY_GNU_VERSION=$VERSION
}

# Discover the Cray MPICH version loaded in Cray
# Expected environment variables are:
# $ echo CRAY_MPICH_VER
# $ echo CRAY_MPICH_VERSION
# $ echo CRAY_MPICH2_VER
function discover_cray_mpich_version(){
  local VERSION=""

  if [[ -n "$CRAY_MPICH_VER" ]]; then
    VERSION=${CRAY_MPICH_VER:0:1}
  elif [[ -n "$CRAY_MPICH_VERSION" ]]; then
    VERSION=${CRAY_MPICH_VERSION:0:1}
  elif [[ -n "$CRAY_MPICH2_VER" ]]; then
    VERSION=${CRAY_MPICH2_VER:0:1}
  else
    echo "Cannot determine Cray MPICH version"
    exit 1
  fi

  SUPPORTED_CRAY_MPICH=(7 8) # We only have wrappers for these versions
  if [[ ! "${SUPPORTED_CRAY_MPICH[@]}" =~ "${VERSION}" ]]; then
    echo "Unsupported CRAY_MPICH version, version $VERSION not supported"
    exit 1
  fi

  if [[ "${VERSION}" -gt "7" ]]; then
    ADD_CRAY_MPICH_VERSION="-$VERSION"
  fi
  CRAY_MPICH_MAJOR_VERSION=$VERSION
}

# Cray Clang is unsupported for Cray MPICH 7.
function check_mpich_crayclang(){
  if [[ "${CRAY_MPICH_MAJOR_VERSION}" -eq "7" && -n "$CRAY_PE_USE_CLANG" ]]; then
    echo "Unsupported MPI, Cray Clang MPICH not supported"
    exit 1
  fi
}

# Selects the most appropriate wrapper that we currently have
function copywrapper()
{
  echo "Selecting proper wrapper for $1"
  ADD_CRAY_MPICH_VERSION=""
  CRAY_MPICH_MAJOR_VERSION=""

  case "$1" in
        cray-prgenvcray)
                discover_cray_mpich_version
                check_mpich_crayclang
                renamecopy "mpich3-cray${ADD_CRAY_MPICH_VERSION}-prgenvcray-64"
                ;;
        cray-prgenvgnu)
                discover_cray_mpich_version
                CRAY_GNU_VERSION=""
                discover_cray_gnu_version
                renamecopy "mpich3-cray${ADD_CRAY_MPICH_VERSION}-prgenvgnu${CRAY_GNU_VERSION}-64"
                ;;
        cray-prgenvintel)
                discover_cray_mpich_version
                renamecopy "mpich3-cray${ADD_CRAY_MPICH_VERSION}-prgenvintel-64"
                ;;
        cray-prgenvpgi)
                discover_cray_mpich_version
                if [[ -z "${ADD_CRAY_MPICH_VERSION}" ]]; then
                  renamecopy "mpich3-cray-prgenvpgi-64"
                else
                  echo "PGI only supported for Cray MPICH 7."
                  exit 1
                fi
                ;;
        cray-prgenvnvidia)
                discover_cray_mpich_version
                if [[ -n "${ADD_CRAY_MPICH_VERSION}" ]]; then
                  renamecopy "mpich3-cray${ADD_CRAY_MPICH_VERSION}-prgenvnvidia-64"
                else
                    echo "NVIDIA only supported for Cray MPICH 8"
                    exit 1
                fi
                ;;
        mpich-4.3|mpich-4.2.3)
                renamecopy "mpich423-gnu-64"
                ;;
        mpich-4.2)
                renamecopy "mpich42-gnu-64"
                ;;
        mpich-4.1)
                renamecopy "mpich41-gnu-64"
                ;;
        mpich-4.0)
                renamecopy "mpich4-gnu-64"
                ;;
        mpich-3.4|mpich-3.3|mpich-3.2|mpich-3.1|mpich-3.0|intelmpi-2021|intelmpi-2019|intelmpi-2018|intelmpi-2017|intelmpi-5)
                renamecopy "mpich3-gnu-64"
                ;;
        OpenMPI-4.0|OpenMPI-4.1)
                renamecopy "openmpi40-gnu-64"
                ;;
        OpenMPI-5.0)
                renamecopy "openmpi50-gnu-64"
                ;;
        mpich2|intelmpi-4)
                renamecopy "mpich2-gnu-64"
                ;;
        *)
                echo "Unsupported MPI version $1"
                exit 1
                ;;
  esac
}

# This function is only called for openmpi currently, since the exact version matters when selecting wrappers
function copywrapper_with_version(){

  MPIVERSION=$($MPIEXECUTABLE --version 2>&1 | grep -o -e "[0-9]\+\.[0-9]\+" | head -1)

  MPINAME="$1-$MPIVERSION"

  echo "MPI version found is $MPINAME"

  copywrapper $MPINAME

}

# This function is called for mpich 4 as support version introduced fix.
function copywrapper_mpich_4(){

  MPIVERSION=$($MPIEXECUTABLE --version 2>&1 | grep -o -e "[0-9]\+\.[0-9]\+" | head -1)

  # A crucial f08 fix made it into the release 4.2.3, make sure that this
  # is captured separately.
  if [[ "${MPIVERSION}"  == "4.2" ]]; then
    FULLVERSION=$($MPIEXECUTABLE --version 2>&1 | grep -o -e "[0-9]\+\.[0-9]\+\.[0-9]\+" | head -1)
    MINORVERSION=${FULLVERSION:4}

    if [[ "${MINORVERSION}" -ge "3" ]]; then
      MPIVERSION="4.2.3"
    fi
  fi

  MPINAME="$1-$MPIVERSION"

  echo "MPI version found is $MPINAME"

  copywrapper $MPINAME

}

# Copywrapper used for intelmpi, gets the exact version of intel mpi. From 2017
# onwards version number is given by date of release.
#
# Expected output of mpirun --version:
# - Ver 2021.6:
#   $ mpirun -version
#     Intel(R) MPI Library for Linux* OS, Version 2021.6 Build 20220227 (id: 28877f3f32)
#     Copyright 2003-2022, Intel Corporation.
# - Ver 2017.4:
#   $ mpirun -version
#      Intel(R) MPI Library for Linux* OS, Version 2017 Update 4 Build 20170817 (id: 17752)
#      Copyright (C) 2003-2017, Intel Corporation. All rights reserved.
# - Ver 4.1.0:
#   $ mpirun --version
#      Intel(R) MPI Library for Linux* OS, Version 4.1.0 Build 20130116
#      Copyright (C) 2003-2013, Intel Corporation. All rights reserved.
function copywrapper_intel_mpi(){
  MPIVERSION=$($MPIEXECUTABLE --version 2>&1 | grep -o -e "Version [0-9]\+" | head -1 | sed 's/\Version //g')

  MPINAME="intelmpi-$MPIVERSION"

  echo "MPI version found is $MPINAME"

  copywrapper $MPINAME
}

# If it's a Cray executable, we give it the mpich3 wrapper (which should work)
if [ "$IS_CRAY" = "yes" ]; then
    case "$PE_ENV" in
        CRAY)
            copywrapper "cray-prgenvcray"
            ;;
        GNU)
            copywrapper "cray-prgenvgnu"
            ;;
        INTEL)
            copywrapper "cray-prgenvintel"
            ;;
        PGI)
            copywrapper "cray-prgenvpgi"
            ;;
        NVIDIA)
            copywrapper "cray-prgenvnvidia"
            ;;
        *)
            echo "Cray PrgEnv ${PE_ENV} unrecognised"
            exit 1
            ;;
    esac
fi

echo "Initial processing of MPI names"
case $MPINAME in
        "OpenMPI")
                copywrapper_with_version $MPINAME
        ;;
        "intel-mpi mpmd"|"intel-mpi mpd"|"intel-mpi")
                copywrapper_intel_mpi
        ;;
        "mpich 4 mpmd")
                copywrapper_mpich_4 "mpich"
        ;;
        "mpich 3 mpmd"|"mpich 3")
                copywrapper_with_version "mpich"
        ;;
        "mpich 2 mpmd"|"mpich 2")
                copywrapper "mpich2"
        ;;
        "bullx mpi")
                copywrapper_with_version "bullxmpi"
        ;;
        "mvapich")
                copywrapper_with_version "mpich"
        ;;
        *)
                echo "$MPINAME not currently supported by precompiled wrappers"
                exit 1
        ;;
esac
