#!/usr/bin/env bash
# shellcheck disable=SC1090
# shellcheck disable=SC1091
set -ex

declare -gx FILES_LIB="./common/handle-files"
if ! [[ -f "${FILES_LIB}" ]] ; then
  echo "error: files-lib is not available"
  exit 1
else
  source "${FILES_LIB}"
fi

declare -agx filepaths=(
  "${PREFIX}/include/cudensitymat.h"
  "${PREFIX}/lib/libcudensitymat.so"
)

if ! files exist "${filepaths[@]}" ; then
  echo "error: cudensitymat files are not available"
  exit 1
fi


declare -gx LINK_MPI=""
declare -gx MPI_LIBRARY EXTRA_LIBS=""
if [[ -f "${PREFIX}/lib/libmpi.so" ]] ; then
  MPI_LIBRARY="${PREFIX}/lib/libmpi.so"
  EXTRA_LIBS="${PREFIX}/lib/libmpi.so"
  LINK_MPI="-lmpi"
fi

declare -gx LINK_NCCL=""
if ls "${PREFIX}/lib/libnccl.so"* >/dev/null 2>&1 ; then
  LINK_NCCL="-lnccl"
fi

declare -gx CUQUANTUM_VER="${1}"
if [[ -n "${CUQUANTUM_VER}" ]] ; then
  declare -agx CUQUANTUM_VERSION
  mapfile -t CUQUANTUM_VERSION <<< "${CUQUANTUM_VER//./$'\n'}"
  {(

    declare -gx IFS_PREV="${IFS}" ; IFS="."
      git clone --branch "v${CUQUANTUM_VERSION[*]:0:3}" "https://github.com/NVIDIA/cuQuantum.git" "cuquantum-cudensitymat"
    IFS="${IFS_PREV}" ; unset IFS_PREV
    cd "cuquantum-cudensitymat/samples/cudensitymat" || exit 1
    declare -agx examples
    mapfile -t examples < <(
      find . -name "*.cpp" -type f -print
    )
    declare -a compile_command
    export LD_LIBRARY_PATH="${PREFIX}/lib"
    declare example ; for example in "${examples[@]}" ; do
      if [[ "${example}" != *"mpi"* && "${example}" != *"nccl"* ]] ; then
        echo "${example}"
        compile_command=(
          "${NVCC:-"nvcc"}" "${NVCC_FLAGS:-"--verbose"}" "--std=c++17"
          "-I${PREFIX}/include" "-L${PREFIX}/lib"
          "-lcudensitymat" "-lcutensornet" "-lcutensor"
          "${example}" "-o" "${example}.out"
        )
        echo "${compile_command[@]}"
        "${compile_command[@]}"
      fi
    done

if [[ "${mpi}" != "nompi" && "${mpi_type}" != "external" && -n "${LINK_MPI}" ]] ; then
      declare example ; for example in "${examples[@]}" ; do
        if [[ "${example}" == *"mpi"* && "${example}" != *"nccl"* ]] ; then
          echo "${example}"
          compile_command=(
            "${NVCC:-"nvcc"}" "${NVCC_FLAGS:-"--verbose"}" "--std=c++17"
            "-I${PREFIX}/include" "-L${PREFIX}/lib"
            "-lcudensitymat" "-lcutensornet" "-lcutensor" "${LINK_MPI}"
            "${example}" "-o" "${example}.out"
          )
          echo "${compile_command[@]}"
          "${compile_command[@]}"
        fi
        if [[ "${example}" == *"nccl"* ]] ; then
          echo "${example}"
          compile_command=(
            "${NVCC:-"nvcc"}" "${NVCC_FLAGS:-"--verbose"}" "--std=c++17"
            "-I${PREFIX}/include" "-L${PREFIX}/lib"
            "-lcudensitymat" "-lcutensornet" "-lcutensor" "${LINK_NCCL}"
            "${example}" "-o" "${example}.out"
          )
          echo "${compile_command[@]}"
          "${compile_command[@]}"
        fi
      done
elif [[ "${mpi}" != "nompi" && "${mpi_type}" == "external" ]] ; then
  echo "external mpi variant; skipping mpi examples"
    fi
  )}
fi

# check for existence of distributed interface library if this is an MPI build
if [[ "${mpi}" != "nompi" && "${mpi_type}" != "external" && -f "${PREFIX}/lib/libmpi.so" ]] ; then
  if ! [[ -f "${CUDENSITYMAT_COMM_LIB}" ]]; then
    echo "error: cudensitymat distributed interface not set in an MPI build"
    exit 1
  fi
  if ! [[ -f "${PREFIX}/lib/libcudensitymat_distributed_interface_nccl.so" ]]; then
    echo "error: cudensitymat nccl distributed interface not found in an MPI build"
    exit 1
  fi
fi