#!/bin/bash -e

# TODO: Review and if possible fix shellcheck errors.
# shellcheck disable=all

[ "${BASH_SOURCE[0]}" ] && SCRIPT_NAME="${BASH_SOURCE[0]}" || SCRIPT_NAME=$0
SCRIPT_DIR="$(cd "$(dirname "$SCRIPT_NAME")/.." && pwd -P)"

spfft_ver="1.1.1"
spfft_sha256="675a048124a96b8c7f89d59d3ac0355833e28b38622e76c4d478ee91b25d766c"
source "${SCRIPT_DIR}"/common_vars.sh
source "${SCRIPT_DIR}"/tool_kit.sh
source "${SCRIPT_DIR}"/signal_trap.sh
source "${INSTALLDIR}"/toolchain.conf
source "${INSTALLDIR}"/toolchain.env

[ -f "${BUILDDIR}/setup_SpFFT" ] && rm "${BUILDDIR}/setup_SpFFT"

! [ -d "${BUILDDIR}" ] && mkdir -p "${BUILDDIR}"
cd "${BUILDDIR}"

case "${with_spfft}" in
  __INSTALL__)
    echo "==================== Installing SpFFT ===================="
    pkg_install_dir="${INSTALLDIR}/SpFFT-${spfft_ver}"
    install_lock_file="${pkg_install_dir}/install_successful"
    if verify_checksums "${install_lock_file}"; then
      echo "SpFFT-${spfft_ver} is already installed, skipping it."
    else
      retrieve_package "${spfft_sha256}" "SpFFT-${spfft_ver}.tar.gz"
      echo "Installing from scratch into ${pkg_install_dir}"
      if [ "${MATH_MODE}" = "mkl" ]; then
        EXTRA_CMAKE_FLAGS="-DSPFFT_MKL=ON -DSPFFT_FFTW_LIB=MKL"
      else
        EXTRA_CMAKE_FLAGS="-DSPFFT_MKL=OFF -DSPFFT_FFTW_LIB=FFTW"
      fi

      [ -d SpFFT-${spfft_ver} ] && rm -rf SpFFT-${spfft_ver}
      tar -xzf SpFFT-${spfft_ver}.tar.gz
      cd SpFFT-${spfft_ver}
      mkdir build-cpu
      cd build-cpu
      cmake \
        -DCMAKE_INSTALL_PREFIX="${pkg_install_dir}" \
        -DCMAKE_INSTALL_LIBDIR=lib \
        -DCMAKE_CXX_COMPILER="${MPICXX}" \
        -DCMAKE_VERBOSE_MAKEFILE=ON \
        -DCMAKE_BUILD_TYPE="RelWithDebInfo" \
        -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
        -DSPFFT_OMP=ON \
        -DSPFFT_MPI=ON \
        -DSPFFT_STATIC=ON \
        -DSPFFT_FORTRAN=ON \
        -DSPFFT_INSTALL=ON \
        ${EXTRA_CMAKE_FLAGS} .. \
        > cmake.log 2>&1 || tail_excerpt cmake.log
      make -j $(get_nprocs) install > make.log 2>&1 || tail_excerpt make.log
      cd ..

      if [ "$ENABLE_CUDA" = "__TRUE__" ]; then
        [ -d build-cuda ] && rm -rf "build-cuda"
        mkdir build-cuda
        cd build-cuda
        cmake \
          -DCMAKE_INSTALL_PREFIX="${pkg_install_dir}" \
          -DCMAKE_INSTALL_LIBDIR=lib \
          -DCMAKE_CXX_COMPILER="${MPICXX}" \
          -DCMAKE_CUDA_FLAGS="-std=c++14 -allow-unsupported-compiler" \
          -DCMAKE_VERBOSE_MAKEFILE=ON \
          -DCMAKE_BUILD_TYPE="RelWithDebInfo" \
          -DSPFFT_OMP=ON \
          -DSPFFT_MPI=ON \
          -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
          -DSPFFT_STATIC=ON \
          -DSPFFT_FORTRAN=ON \
          -DSPFFT_INSTALL=ON \
          -DSPFFT_GPU_BACKEND=CUDA \
          ${EXTRA_CMAKE_FLAGS} .. > cmake.log 2>&1 || tail_excerpt cmake.log
        make -j $(get_nprocs) > make.log 2>&1 || tail_excerpt make.log
        install -d ${pkg_install_dir}/lib/cuda
        [ -f src/libspfft.a ] && install -m 644 src/*.a ${pkg_install_dir}/lib/cuda >> install.log 2>&1
        [ -f src/libspfft.so ] && install -m 644 src/*.so ${pkg_install_dir}/lib/cuda >> install.log 2>&1
      fi

      if [ "$ENABLE_HIP" = "__TRUE__" ]; then
        case "${GPUVER}" in
          K20X | K40 | K80 | P100 | V100 | A100 | A40 | H100)
            [ -d build-cuda ] && rm -rf "build-cuda"
            mkdir build-cuda
            cd build-cuda
            cmake \
              -DCMAKE_INSTALL_PREFIX="${pkg_install_dir}" \
              -DCMAKE_INSTALL_LIBDIR=lib \
              -DCMAKE_CXX_COMPILER="${MPICXX}" \
              -DCMAKE_CUDA_FLAGS="-std=c++14 -allow-unsupported-compiler" \
              -DCMAKE_VERBOSE_MAKEFILE=ON \
              -DCMAKE_BUILD_TYPE="RelWithDebInfo" \
              -DSPFFT_OMP=ON \
              -DSPFFT_MPI=ON \
              -DSPFFT_STATIC=ON \
              -DSPFFT_FORTRAN=ON \
              -DSPFFT_INSTALL=ON \
              -DSPFFT_GPU_BACKEND=CUDA \
              ${EXTRA_CMAKE_FLAGS} .. > cmake.log 2>&1 || tail_excerpt cmake.log
            make -j $(get_nprocs) > make.log 2>&1 || tail_excerpt make.log
            install -d ${pkg_install_dir}/lib/cuda
            [ -f src/libspfft.a ] && install -m 644 src/*.a ${pkg_install_dir}/lib/cuda >> install.log 2>&1
            [ -f src/libspfft.so ] && install -m 644 src/*.so ${pkg_install_dir}/lib/cuda >> install.log 2>&1
            ;;
          Mi50 | Mi100 | Mi200 | Mi250)
            [ -d build-hip ] && rm -rf "build-hip"
            mkdir build-hip
            cd build-hip
            cmake \
              -DCMAKE_INSTALL_PREFIX="${pkg_install_dir}" \
              -DCMAKE_INSTALL_LIBDIR=lib \
              -DCMAKE_CXX_COMPILER="${MPICXX}" \
              -DCMAKE_VERBOSE_MAKEFILE=ON \
              -DCMAKE_BUILD_TYPE="RelWithDebInfo" \
              -DSPFFT_OMP=ON \
              -DSPFFT_MPI=ON \
              -DSPFFT_STATIC=ON \
              -DSPFFT_FORTRAN=ON \
              -DSPFFT_INSTALL=ON \
              -DSPFFT_GPU_BACKEND=ROCM \
              ${EXTRA_CMAKE_FLAGS} .. \
              > cmake.log 2>&1 || tail_excerpt cmake.log
            make -j $(get_nprocs) > make.log 2>&1 || tail_excerpt make.log
            install -d ${pkg_install_dir}/lib/rocm
            [ -f src/libspfft.a ] && install -m 644 src/*.a ${pkg_install_dir}/lib/rocm >> install.log 2>&1
            [ -f src/libspfft.so ] && install -m 644 src/*.so ${pkg_install_dir}/lib/rocm >> install.log 2>&1
            ;;
          *) ;;

        esac
      fi
      write_checksums "${install_lock_file}" "${SCRIPT_DIR}/stage8/$(basename ${SCRIPT_NAME})"
    fi
    ;;
  __SYSTEM__)
    echo "==================== Finding SpFFT from system paths ===================="
    check_pkgconfig spfft
    pkg_install_dir="$(pkg-config --variable=prefix spfft)"
    ;;
  __DONTUSE__)
    # Nothing to do
    ;;
  *)
    echo "==================== Linking SpFFT to user paths ===================="
    pkg_install_dir="${with_spfft}"
    # use the lib64 directory if present (multi-abi distros may link lib/ to lib32/ instead)
    SPFFT_LIBDIR="${pkg_install_dir}/lib"
    [ -d "${pkg_install_dir}/lib64" ] && SPFFT_LIBDIR="${pkg_install_dir}/lib64"
    check_dir "${SPFFT_LIBDIR}"
    check_dir "${pkg_install_dir}/include"
    ;;
esac
if [ "$with_spfft" != "__DONTUSE__" ]; then
  if [ "$with_spfft" != "__SYSTEM__" ]; then
    cat << EOF > "${BUILDDIR}/setup_spfft"
prepend_path LD_LIBRARY_PATH "${pkg_install_dir}/lib"
prepend_path LD_RUN_PATH "${pkg_install_dir}/lib"
prepend_path LIBRARY_PATH "${pkg_install_dir}/lib"
prepend_path PKG_CONFIG_PATH "${pkg_install_dir}/lib/pkgconfig"
prepend_path CMAKE_PREFIX_PATH "${pkg_install_dir}"
EOF
  fi
  cat << EOF >> "${BUILDDIR}/setup_spfft"
export SpFFT_ROOT="${pkg_install_dir}"
export SPFFT_VER="${spfft_ver}"
EOF
  filter_setup "${BUILDDIR}/setup_spfft" "${SETUPFILE}"
fi

load "${BUILDDIR}/setup_spfft"
write_toolchain_env "${INSTALLDIR}"

cd "${ROOTDIR}"
report_timing "spfft"
