#!/bin/bash
mydir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"

export HPCX_DIR=$mydir
export HPCX_UCX_DIR=${HPCX_DIR}/ucx/mt
export HPCX_UCC_DIR=${HPCX_DIR}/ucc
export HPCX_SHARP_DIR=${HPCX_DIR}/sharp
export HPCX_HCOLL_DIR=${HPCX_DIR}/hcoll
export HPCX_MPI_DIR=${HPCX_DIR}/ompi
export HPCX_OSHMEM_DIR=${HPCX_MPI_DIR}
export HPCX_MPI_TESTS_DIR=${HPCX_MPI_DIR}/tests
export HPCX_OSU_DIR=${HPCX_MPI_TESTS_DIR}/osu-micro-benchmarks
export HPCX_OSU_CUDA_DIR=${HPCX_MPI_TESTS_DIR}/osu-micro-benchmarks-cuda
export HPCX_CLUSTERKIT_DIR=${HPCX_DIR}/clusterkit
export HPCX_NCCLNET_PLUGIN_DIR=${HPCX_DIR}/nccl_spectrum-x_plugin
if [[ "${HPCX_ENABLE_NCCL_MRC_PLUGIN}" == "1" ]]; then
    export HPCX_NCCLNET_PLUGIN_DIR=${HPCX_DIR}/nccl_mrc_plugin
fi
export HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR=${HPCX_DIR}/nccl_rdma_sharp_plugin
export OMPI_HOME=${HPCX_MPI_DIR}
export MPI_HOME=${HPCX_MPI_DIR}
export OSHMEM_HOME=${HPCX_MPI_DIR}
export SHMEM_HOME=${HPCX_MPI_DIR}

function hpcx_load() {
    PATH=${PATH:-""}
    OPAL_PREFIX=${OPAL_PREFIX:-""}
    export OLD_PATH=$PATH
    export OLD_OPAL_PREFIX=${OPAL_PREFIX}
    export OPAL_PREFIX=${HPCX_MPI_DIR}
    export PATH=${HPCX_MPI_DIR}/bin${PATH:+:$PATH}
    export PATH=${HPCX_UCX_DIR}/bin${PATH:+:$PATH}
    export PATH=${HPCX_UCC_DIR}/bin${PATH:+:$PATH}
    export PATH=${HPCX_HCOLL_DIR}/bin${PATH:+:$PATH}
    export PATH=${HPCX_CLUSTERKIT_DIR}/bin${PATH:+:$PATH}
    export PATH=${HPCX_SHARP_DIR}/bin${PATH:+:$PATH}

    LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-""}
    export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_MPI_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_HCOLL_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_SHARP_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_UCX_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_UCX_DIR}/lib/ucx${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_UCC_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    export LD_LIBRARY_PATH=${HPCX_UCC_DIR}/lib/ucc${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}

    if [[ "${HPCX_ENABLE_NCCL_SPCX_PLUGIN}" == "1" || "${HPCX_ENABLE_NCCL_MRC_PLUGIN}" == "1" ]]; then
        export LD_LIBRARY_PATH=${HPCX_NCCLNET_PLUGIN_DIR}/lib:${HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    else
        export LD_LIBRARY_PATH=${HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR}/lib:${HPCX_NCCLNET_PLUGIN_DIR}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
    fi

    LIBRARY_PATH=${LIBRARY_PATH:-""}
    export OLD_LIBRARY_PATH=${LIBRARY_PATH}
    export LIBRARY_PATH=${HPCX_MPI_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}
    export LIBRARY_PATH=${HPCX_HCOLL_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}
    export LIBRARY_PATH=${HPCX_UCX_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}
    export LIBRARY_PATH=${HPCX_UCC_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}
    export LIBRARY_PATH=${HPCX_SHARP_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}

    if [[ "${HPCX_ENABLE_NCCL_SPCX_PLUGIN}" == "1" || "${HPCX_ENABLE_NCCL_MRC_PLUGIN}" == "1" ]]; then
        export LIBRARY_PATH=${HPCX_NCCLNET_PLUGIN_DIR}/lib:${HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}
    else
        export LIBRARY_PATH=${HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR}/lib:${HPCX_NCCLNET_PLUGIN_DIR}/lib${LIBRARY_PATH:+:$LIBRARY_PATH}
    fi

    CPATH=${CPATH:-""}
    export OLD_CPATH=$CPATH
    export CPATH=${HPCX_HCOLL_DIR}/include${CPATH:+:$CPATH}
    export CPATH=${HPCX_SHARP_DIR}/include${CPATH:+:$CPATH}
    export CPATH=${HPCX_UCX_DIR}/include${CPATH:+:$CPATH}
    export CPATH=${HPCX_UCC_DIR}/include${CPATH:+:$CPATH}
    export CPATH=${HPCX_MPI_DIR}/include${CPATH:+:$CPATH}

    PKG_CONFIG_PATH=${PKG_CONFIG_PATH:-""}
    export OLD_PKG_CONFIG_PATH=${PKG_CONFIG_PATH}
    export PKG_CONFIG_PATH=${HPCX_MPI_DIR}/lib/pkgconfig${PKG_CONFIG_PATH:+:$PKG_CONFIG_PATH}
    export PKG_CONFIG_PATH=${HPCX_HCOLL_DIR}/lib/pkgconfig:${HPCX_SHARP_DIR}/lib/pkgconfig:${HPCX_UCX_DIR}/lib/pkgconfig${PKG_CONFIG_PATH:+:$PKG_CONFIG_PATH}
}

function hpcx_unload() {
    if [[ -z "${OLD_PATH+x}" ]]; then
        return 0
    fi

    export PATH="${OLD_PATH}"
    export LD_LIBRARY_PATH="${OLD_LD_LIBRARY_PATH}"
    export LIBRARY_PATH="${OLD_LIBRARY_PATH}"
    export OPAL_PREFIX="${OLD_OPAL_PREFIX}"
    export CPATH="${OLD_CPATH}"
    export PKG_CONFIG_PATH="${OLD_PKG_CONFIG_PATH}"

    unset OLD_PATH
    unset OLD_LD_LIBRARY_PATH
    unset OLD_LIBRARY_PATH
    unset OLD_OPAL_PREFIX
    unset OLD_CPATH
    unset OLD_PKG_CONFIG_PATH

    for var in $(env | cut -f1 -d= | grep '^HPCX_'); do
        unset "$var"
    done

    for var in MPI_HOME OSHMEM_HOME SHMEM_HOME OMPI_HOME; do
        unset "$var"
    done
}
