{ lib
, fetchFromGitHub
, cmake
, withCFlags
, intelPackages

, mklSupport ? true

, config

, cudaSupport ? config.cudaSupport
, cudaPackages ? { }

, rocmSupport ? config.rocmSupport
, hipTargets ? null # only one target at a time supported
, rocmPackages ? { }
}:

let
  # rocmSupport is not enough, we need a specific target
  enableHip = rocmSupport && hipTargets != null;

  stdenv = withCFlags (lib.optionals cudaSupport [ "--cuda-path=${cudaPackages.cudatoolkit}" ]) intelPackages.stdenv;
in

  # at least one backend has to be enabled
  assert mklSupport || cudaSupport || enableHip;

stdenv.mkDerivation rec {
  pname = "oneMath";
  version = "0.8";

  src = fetchFromGitHub {
    owner = "uxlfoundation";
    repo = "oneMath";
    rev = "v${version}";
    sha256 = "sha256-xK8lKI3oqKlx3xtvdScpMq+HXAuoYCP0BZdkEqnJP5o=";
  };

  cmakeFlags = [
    (lib.cmakeBool "ENABLE_MKLCPU_BACKEND"    mklSupport)
    (lib.cmakeBool "ENABLE_MKLGPU_BACKEND"    mklSupport)

    (lib.cmakeBool "ENABLE_CUBLAS_BACKEND"    cudaSupport)
    (lib.cmakeBool "ENABLE_CUFFT_BACKEND"     cudaSupport)
    (lib.cmakeBool "ENABLE_CURAND_BACKEND"    cudaSupport)
    (lib.cmakeBool "ENABLE_CUSOLVER_BACKEND"  cudaSupport)
    (lib.cmakeBool "ENABLE_CUSPARSE_BACKEND"  cudaSupport)

    (lib.cmakeBool "ENABLE_ROCBLAS_BACKEND"   enableHip)
    (lib.cmakeBool "ENABLE_ROCFFT_BACKEND"    enableHip)
    (lib.cmakeBool "ENABLE_ROCSOLVER_BACKEND" enableHip)
    (lib.cmakeBool "ENABLE_ROCRAND_BACKEND"   enableHip)
    (lib.cmakeBool "ENABLE_ROCSPARSE_BACKEND" enableHip)

    (lib.cmakeBool "BUILD_FUNCTIONAL_TESTS" false)
    (lib.cmakeBool "BUILD_EXAMPLES" false)
  ] ++ lib.optionals enableHip [
    (lib.cmakeFeature "HIP_TARGETS" hipTargets)
  ];

  nativeBuildInputs = [ cmake ];
  buildInputs = lib.optionals (mklSupport) [
    intelPackages.mkl
    intelPackages.tbb
  ] ++ lib.optionals (enableHip) [
    rocmPackages.rocmPath
    rocmPackages.rocblas
    rocmPackages.rocfft
    rocmPackages.rocsolver
    rocmPackages.rocrand
    rocmPackages.rocsparse
  ] ++ lib.optionals (cudaSupport) [
    (lib.getDev cudaPackages.cuda_cudart)
    cudaPackages.cudatoolkit

    cudaPackages.libcublas
    cudaPackages.libcurand
    cudaPackages.libcufft
    cudaPackages.libcusparse
    cudaPackages.libcusolver
  ];

}
