#ifndef LIBRAPID_MUTLIARRAY_OPERATIONS
#define LIBRAPID_MUTLIARRAY_OPERATIONS

#include <librapid/config.hpp>
#include <librapid/utils/time_utils.hpp>
#include <librapid/autocast/autocast.hpp>
#include <librapid/config.hpp>
#include <librapid/array/stride.hpp>
#include <librapid/array/extent.hpp>

namespace librapid::imp {
#ifdef LIBRAPID_HAS_CUDA
	// CUDA_INCLUDE_DIRS
	inline const jitify::detail::vector<std::string> cudaHeaders = {
	  CUDA_INCLUDE_DIRS + std::string("/curand.h"),
	  CUDA_INCLUDE_DIRS + std::string("/curand_kernel.h"),
	  CUDA_INCLUDE_DIRS + std::string("/cublas_v2.h"),
	  CUDA_INCLUDE_DIRS + std::string("/cublas_api.h"),
	  CUDA_INCLUDE_DIRS + std::string("/cuda_fp16.h"),
	  CUDA_INCLUDE_DIRS + std::string("/cuda_bf16.h")};

	inline const std::vector<std::string> cudaParams = {"--disable-warnings",
														"-std=c++17",
														std::string("-I") +
														  CUDA_INCLUDE_DIRS};

	inline constexpr char complexHpp[] = R"V0G0N(
#ifndef LIBRAPID_CUSTOM_COMPLEX
#define LIBRAPID_CUSTOM_COMPLEX
namespace librapid
{
	template <class T>
	class Complex
	{
	public:
		Complex(const T &real_val = T(), const T &imag_val = T())
			: m_real(real_val), m_imag(imag_val)
		{}
		Complex &operator=(const T &val)
		{
			m_real = val;
			m_imag = 0;
			return *this;
		}
		template <class V>
		Complex(const Complex<V> &other)
			: Complex(static_cast<T>(other.real()), static_cast<T>(other.imag()))
		{}
		template <class V>
		Complex &operator=(const Complex<V> &other)
		{
			m_real = static_cast<T>(other.real());
			m_imag = static_cast<T>(other.imag());
			return *this;
		}
		Complex copy() const
		{
			return Complex<T>(m_real, m_imag);
		}
		Complex operator-() const
		{
			return Complex<T>(-m_real, -m_imag);
		}
		template<typename V>
		Complex operator+(const V &other) const
		{
			return Complex<T>(m_real + other, m_imag);
		}
		template<typename V>
		Complex operator-(const V &other) const
		{
			return Complex<T>(m_real - other, m_imag);
		}
		template<typename V>
		Complex operator*(const V &other) const
		{
			return Complex<T>(m_real * other, m_imag * other);
		}
		template<typename V>
		Complex operator/(const V &other) const
		{
			return Complex<T>(m_real / other, m_imag / other);
		}
		template<typename V>
		Complex &operator+=(const V &other)
		{
			m_real += other;
			return *this;
		}
		template<typename V>
		Complex &operator-=(const V &other)
		{
			m_real -= other;
			return *this;
		}
		template<typename V>
		Complex &operator*=(const V &other)
		{
			m_real *= other;
			m_imag *= other;
			return *this;
		}
		template<typename V>
		Complex &operator/=(const V &other)
		{
			m_real /= other;
			m_imag /= other;
			return *this;
		}
		template<typename V>
		Complex operator+(const Complex<V> &other) const
		{
			return Complex(m_real + other.real(),
						   m_imag + other.imag());
		}
		template<typename V>
		Complex operator-(const Complex<V> &other) const
		{
			return Complex(m_real - other.real(),
						   m_imag - other.imag());
		}
		template<typename V>
		Complex operator*(const Complex<V> &other) const
		{
			return Complex((m_real * other.real()) - (m_imag * other.imag()),
						   (m_real * other.imag()) + (m_imag * other.real()));
		}
		template<typename V>
		Complex operator/(const Complex<V> &other) const
		{
			return Complex((m_real * other.real()) + (m_imag * other.imag()) /
						   ((other.real() * other.real()) + (other.imag() * other.imag())),
						   (m_real * other.real()) - (m_imag * other.imag()) /
						   ((other.real() * other.real()) + (other.imag() * other.imag())));
		}
		template<typename V>
		Complex &operator+=(const Complex<V> &other)
		{
			m_real = m_real + other.real();
			m_imag = m_imag + other.imag();
			return *this;
		}
		template<typename V>
		Complex &operator-=(const Complex<V> &other)
		{
			m_real = m_real - other.real();
			m_imag = m_imag - other.imag();
			return *this;
		}
		template<typename V>
		Complex &operator*=(const Complex<V> &other)
		{
			m_real = (m_real * other.real()) - (m_imag * other.imag());
			m_imag = (m_real * other.imag()) + (imag() * other.real());
			return *this;
		}
		template<typename V>
		Complex &operator/=(const Complex<V> &other)
		{
			m_real = (m_real * other.real()) + (m_imag * other.imag()) /
				((other.real() * other.real()) + (other.imag() * other.imag()));
			m_imag = (m_real * other.real()) - (m_imag * other.imag()) /
				((other.real() * other.real()) + (other.imag() * other.imag()));
			return *this;
		}
		template<typename V>
		bool operator==(const Complex<V> &other) const
		{
			return m_real == other.real() && m_imag == other.imag();
		}
		template<typename V>
		bool operator!=(const Complex<V> &other) const
		{
			return !(*this == other);
		}
		T mag() const
		{
			return std::sqrt(m_real * m_real + m_imag * m_imag);
		}
		T angle() const
		{
			return std::atan2(m_real, m_imag);
		}
		Complex<T> log() const
		{
			return Complex<T>(std::log(mag()), angle());
		}
		Complex<T> conjugate() const
		{
			return Complex<T>(m_real, -m_imag);
		}
		Complex<T> reciprocal() const
		{
			return Complex<T>((m_real) / (m_real * m_real + m_imag * m_imag),
							  -(m_imag) / (m_real * m_real + m_imag * m_imag));
		}
		const T &real() const
		{
			return m_real;
		}
		T &real()
		{
			return m_real;
		}
		const T &imag() const
		{
			return m_imag;
		}
		T &imag()
		{
			return m_imag;
		}
		template<typename V>
		operator V() const
		{
			return m_real;
		}
	private:
		T m_real = 0;
		T m_imag = 0;
	};
	template<typename A, typename B>
		Complex<B> operator+(const A &a, const Complex<B> &b)
	{
		return Complex<B>(a) + b;
	}
	template<typename A, typename B>
		Complex<B> operator-(const A &a, const Complex<B> &b)
	{
		return Complex<B>(a) - b;
	}
	template<typename A, typename B>
		Complex<B> operator*(const A &a, const Complex<B> &b)
	{
		return Complex<B>(a) * b;
	}
	template<typename A, typename B>
		Complex<B> operator/(const A &a, const Complex<B> &b)
	{
		return Complex<B>(a) / b;
	}
	template<typename A, typename B>
		A &operator+=(A &a, const Complex<B> &b)
	{
		a += b.real();
		return a;
	}
	template<typename A, typename B>
		A &operator-=(A &a, const Complex<B> &b)
	{
		a -= b.real();
		return a;
	}
	template<typename A, typename B>
		A &operator*=(A &a, const Complex<B> &b)
	{
		a *= b.real();
		return a;
	}
	template<typename A, typename B>
		A &operator/=(A &a, const Complex<B> &b)
	{
		a /= b.real();
		return a;
	}
}
#endif
		)V0G0N";
#endif // LIBRAPID_HAS_CUDA

	inline int makeSameAccelerator(RawArray &dst, const RawArray &src,
								   int64_t size) {
		// Freeing information
		// 0 = no free
		// 1 = cudaFree()
		// 2 = free()

		auto freeMode = -1;
		if (dst.location == src.location) {
			freeMode = 0;
			dst		 = src;
		} else {
			if (src.location == Accelerator::CPU) {
				// Copy from CPU to GPU
				// Allocate memory

				freeMode = 1;
				rawArrayMalloc(dst, size);
				rawArrayMemcpy(dst, src, size);
			}
#ifdef LIBRAPID_HAS_CUDA
			else if (src.location == Accelerator::GPU) {
				// Copy A from GPU to CPU

				freeMode = 2;
				rawArrayMalloc(dst, size);
				rawArrayMemcpy(dst, src, size);
			}
#else
			else {
				throw std::invalid_argument(
				  "GPU support was not enabled, so"
				  " calculations involving the GPU"
				  " are not possible");
			}
#endif // LIBRAPID_HAS_CUDA
		}

		return freeMode;
	}

	/**
	 * \rst
	 *
	 * Free memory with a given freeing mode
	 *
	 * - 0 = do not free memory
	 * - 1 = free host memory (``free()``)
	 * - 2 = free device memory (``cudaFree()``)
	 *
	 * \endrst
	 */
	inline void freeWithMode(RawArray &raw, int mode) {
		// Freeing information
		// 0 = no free
		// 1 = cudaFree()
		// 2 = free()
		if (mode == 0) return;

		if (mode == 1 || mode == 2)
			rawArrayFree(raw);
		else
			throw std::invalid_argument(
			  "Invalid free mode for binary operation");
	}

	/**
	 * \rst
	 *
	 * Perform a unary operation on trivially strided, contiguous memory.
	 *
	 * An operation struct can be passed to the ``op`` parameter of this
	 * function. It must contain a ``func``, being a lambda taking a single
	 * templated value, and returning a single value of the same type.
	 *
	 * If CUDA support is enabled, the operation struct must also contain a
	 * small kernel operation stored as a string. An example kernel is
	 * ``"b = 2 * a"`` -- "b" and "a" are scalar values corresponding to the
	 * provided input data. This example kernel will multiply every value in
	 * ``a`` by two, and will store the result in ``b``.
	 *
	 *
	 * Parameters
	 * ----------
	 * locnA: Accelerator
	 *		The location of ``a``
	 * locnB: Accelerator
	 *		The location of ``b``
	 * a: ``A *``
	 *		Trivial array
	 * b: ``B *``
	 *		Trivial array
	 * size: integer
	 *		The number of *elements* in ``a`` (and therefore ``b``)
	 * op: Operation struct (containing "name", "kernel" and "func")
	 *
	 * \endrst
	 */
	template<typename FUNC>
	inline void multiarrayUnaryOpTrivial(RawArray dst, const RawArray &src,
										 int64_t elems, const FUNC &op) {
		if (dst.location != src.location) {
			// Copy A to be on the same accelerator as B
			RawArray tempSrc = {(int64_t *)nullptr, dst.dtype, dst.location};
			rawArrayMalloc(tempSrc, elems);
			int freeMode = makeSameAccelerator(tempSrc, src, elems);
			multiarrayUnaryOpTrivial(dst, tempSrc, elems, op);
			freeWithMode(tempSrc, freeMode);
		} else {
			if (dst.location == Accelerator::CPU) {
				std::visit(
				  [&](auto *dstData, auto *srcData) {
					  auto tempElems = elems;
					  auto tempOp	 = op;

					  using A =
						typename std::remove_pointer<decltype(dstData)>::type;
					  using B =
						typename std::remove_pointer<decltype(srcData)>::type;

					  if (elems < THREAD_THREASHOLD) {
						  for (int64_t i = 0; i < tempElems; ++i)
							  dstData[i] = tempOp(srcData[i], i);
					  } else {
#pragma omp parallel for shared(dstData, srcData, tempElems, tempOp)
						  for (int64_t i = 0; i < tempElems; ++i) {
							  dstData[i] = tempOp(srcData[i], i);
						  }
					  }
				  },
				  dst.data,
				  src.data);
			}
#ifdef LIBRAPID_HAS_CUDA
			else {
				using jitify::reflection::Type;

				static double randSeed = seconds() * 10;

				std::string kernel = "unaryKernelTrivial\n";
				kernel += "__constant__ int LIBRAPID_MAX_DIMS = " +
						  std::to_string(LIBRAPID_MAX_DIMS) + ";\n";
				kernel +=
				  "#include <stdint.h>\n"
				  "#include <type_traits>\n"
				  "#include <" CUDA_INCLUDE_DIRS
				  "/curand_kernel.h>\n"
				  "#include <" CUDA_INCLUDE_DIRS "/curand.h>\n\n";

				kernel += complexHpp;

				kernel += R"V0G0N(
					template<typename T_DST, typename A>
					__device__
					inline auto )V0G0N";

				kernel += op.name;
				kernel +=
				  R"V0G0N((const A &a, int64_t indexA, curandState_t *_curandState)
					{
					)V0G0N";
				kernel += op.kernel;

				kernel += R"V0G0N(}
					template<typename T_DST, typename T_SRC>
					__global__
					void unaryFuncTrivial(T_DST *__restrict dstData,
										  const T_SRC *__restrict srcData,
										  const uint64_t size,
										  curandState_t *_curandStates = nullptr)
					{
						int64_t kernelIndex = blockDim.x * blockIdx.x + threadIdx.x;
						if (kernelIndex < size) {
					)V0G0N";
				kernel += "dstData[kernelIndex] = " + op.name +
						  "<T_DST, T_SRC>(srcData[kernelIndex], kernelIndex, "
						  "_curandStates + kernelIndex);";
				kernel += "\n}\n}";

				static jitify::JitCache kernelCache;
				jitify::Program program =
				  kernelCache.program(kernel, cudaHeaders, cudaParams);

				unsigned int threadsPerBlock, blocksPerGrid;

				// Use 1 to 512 threads per block
				if (elems < 512) {
					threadsPerBlock = (unsigned int)elems;
					blocksPerGrid	= 1;
				} else {
					threadsPerBlock = 512;
					blocksPerGrid =
					  ceil(double(elems) / double(threadsPerBlock));
				}

				dim3 grid(blocksPerGrid);
				dim3 block(threadsPerBlock);

				static curandState_t *curandGenerators = nullptr;
				static int64_t numGenerators		   = 0;
				if ((op.name == "fillRandom" ||
					 op.name == "fillRandomComplex") &&
					(numGenerators < threadsPerBlock * blocksPerGrid)) {
					// Free curand generators and reallocate with a larger size
					// Allocate the nearest power of two generators
					cudaSafeCall(cudaFree(curandGenerators));

					numGenerators = 1;
					while (numGenerators < threadsPerBlock * blocksPerGrid)
						numGenerators <<= 1;

#	ifdef LIBRAPID_CUDA_STREAM
					cudaSafeCall(
					  cudaMallocAsync(&curandGenerators,
									  sizeof(curandState_t) * numGenerators,
									  cudaStream));
#	else
					cudaSafeCall(
					  cudaMallocAsync(&curandGenerators,
									  sizeof(curandState_t) * numGenerators));
#	endif // LIBRAPID_CUDA_STREAM

					std::string randomKernel =
					  std::string("randomKernel\n") + "#include <stdint.h>\n" +
					  "#include <type_traits>\n" + "#include <" +
					  CUDA_INCLUDE_DIRS + "/curand_kernel.h>\n" + "#include <" +
					  CUDA_INCLUDE_DIRS + "/curand.h>\n\n";

					randomKernel += R"V0G0N(
						__global__
						void populateGenerators(curandState_t *states, int64_t numStates) {
							int64_t kernelIndex = blockDim.x * blockIdx.x + threadIdx.x;
							curand_init()V0G0N";
					randomKernel +=
					  std::to_string(getSeed(op)) +
					  ", kernelIndex, 0, &states[kernelIndex]);}\n";

					static jitify::JitCache randomCache;
					jitify::Program randomProgram = kernelCache.program(
					  randomKernel, cudaHeaders, cudaParams);

#	ifdef LIBRAPID_CUDA_STREAM
					jitifyCall(randomProgram.kernel("populateGenerators")
								 .instantiate()
								 .configure(grid, block, 0, cudaStream)
								 .launch(curandGenerators, numGenerators));
#	else
					jitifyCall(program.kernel("populateGenerators")
								 .instantiate()
								 .configure(grid, block)
								 .launch(curandGenerators, numGenerators));
#	endif // LIBRAPID_CUDA_STREAM
				}

				std::visit(
				  [&](auto *a, auto *b) {
					  using A = typename std::remove_pointer<decltype(a)>::type;
					  using B = typename std::remove_pointer<decltype(b)>::type;

					  if (op.name == "fillRandom" ||
						  op.name == "fillRandomComplex") {
#	ifdef LIBRAPID_CUDA_STREAM
						  jitifyCall(program.kernel("unaryFuncTrivial")
									   .instantiate(Type<A>(), Type<B>())
									   .configure(grid, block, 0, cudaStream)
									   .launch(a, b, elems, curandGenerators));
#	else
						  jitifyCall(program.kernel("unaryFuncTrivial")
									   .instantiate(Type<A>(), Type<B>())
									   .configure(grid, block)
									   .launch(a, b, elems, curandGenerators));
#	endif // LIBRAPID_CUDA_STREAM
					  } else {
#	ifdef LIBRAPID_CUDA_STREAM
						  jitifyCall(program.kernel("unaryFuncTrivial")
									   .instantiate(Type<A>(), Type<B>())
									   .configure(grid, block, 0, cudaStream)
									   .launch(a, b, elems, nullptr));
#	else
						  jitifyCall(program.kernel("unaryFuncTrivial")
									   .instantiate(Type<A>(), Type<B>())
									   .configure(grid, block)
									   .launch(a, b, elems, nullptr));
#	endif // LIBRAPID_CUDA_STREAM
					  }
				  },
				  dst.data,
				  src.data);
			}
#else
			else {
				throw std::runtime_error(
				  "CUDA support was not enabled. Invalid operation");
			}
#endif // LIBRAPID_HAS_CUDA
		}
	}

	template<typename FUNC>
	inline void
	multiarrayUnaryOpComplex(RawArray dst, RawArray src, int64_t elems,
							 const Extent &extent, const Stride &dstStride,
							 const Stride &srcStride, const FUNC &op,
							 bool trivialDst = false) {
		if (dst.location != src.location) {
			// Copy A to be on the same accelerator as B
			RawArray tempSrc = {(int64_t *)nullptr, dst.dtype, dst.location};
			rawArrayMalloc(tempSrc, elems);
			int freeMode = makeSameAccelerator(tempSrc, src, elems);
			multiarrayUnaryOpTrivial(dst, tempSrc, elems, op);
			freeWithMode(tempSrc, freeMode);
		} else {
			// Locations are equal
			if (dst.location == Accelerator::CPU) {
				// Iterate over the array using its stride and extent

				// Counters
				int64_t idim = 0;
				int64_t ndim = extent.ndim();

				// Create pointers here so repeated function calls
				// are not needed
				static int64_t rawExtent[LIBRAPID_MAX_DIMS];
				static int64_t rawDstStride[LIBRAPID_MAX_DIMS];
				static int64_t rawSrcStride[LIBRAPID_MAX_DIMS];

				for (int64_t i = 0; i < ndim; ++i) {
					rawExtent[ndim - i - 1]	   = extent.raw()[i];
					rawDstStride[ndim - i - 1] = dstStride[i];
					rawSrcStride[ndim - i - 1] = srcStride[i];
				}

				std::visit(
				  [&](auto *dstData, auto *srcData) {
					  using A		   = std::remove_pointer<decltype(dstData)>;
					  using B		   = std::remove_pointer<decltype(srcData)>;
					  int64_t dstIndex = 0, srcIndex = 0;
					  int64_t coord[LIBRAPID_MAX_DIMS] {};

					  do {
						  // *dstData = op(*srcData);
						  dstData[dstIndex] = op(srcData[srcIndex], srcIndex);

						  for (idim = 0; idim < ndim; ++idim) {
							  if (++coord[idim] == rawExtent[idim]) {
								  coord[idim] = 0;
								  // srcData = srcData - (rawExtent[idim] - 1) *
								  // rawSrcStride[idim]; dstData = dstData -
								  // (rawExtent[idim] - 1) * rawDstStride[idim];
								  srcIndex = srcIndex - (rawExtent[idim] - 1) *
														  rawSrcStride[idim];
								  dstIndex = dstIndex - (rawExtent[idim] - 1) *
														  rawDstStride[idim];
							  } else {
								  // srcData = srcData + rawSrcStride[idim];
								  // dstData = dstData + rawDstStride[idim];
								  srcIndex = srcIndex + rawSrcStride[idim];
								  dstIndex = dstIndex + rawDstStride[idim];
								  break;
							  }
						  }
					  } while (idim < ndim);
				  },
				  dst.data,
				  src.data);
			}
#ifdef LIBRAPID_HAS_CUDA
			else {
				using jitify::reflection::Type;

				// Buffers for extent and strides
				int64_t dims = extent.ndim();
				int64_t *deviceExtent, *deviceDstStride = nullptr,
									   *deviceSrcStride;

				// Copy all the required data, including extents and strides
#	ifdef LIBRAPID_CUDA_STREAM
				cudaSafeCall(
				  cudaMallocAsync(&deviceExtent,
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaStream));
				if (!trivialDst)
					cudaSafeCall(
					  cudaMallocAsync(&deviceDstStride,
									  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
									  cudaStream));
				cudaSafeCall(
				  cudaMallocAsync(&deviceSrcStride,
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaStream));

				cudaSafeCall(
				  cudaMemcpyAsync(deviceExtent,
								  extent.raw(),
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaMemcpyHostToDevice,
								  cudaStream));
				if (!trivialDst)
					cudaSafeCall(
					  cudaMemcpyAsync(deviceDstStride,
									  dstStride.raw(),
									  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
									  cudaMemcpyHostToDevice,
									  cudaStream));
				cudaSafeCall(
				  cudaMemcpyAsync(deviceSrcStride,
								  srcStride.raw(),
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaMemcpyHostToDevice,
								  cudaStream));
#	else
				cudaSafeCall(cudaMalloc(&deviceExtent,
										sizeof(int64_t) * LIBRAPID_MAX_DIMS));
				if (!trivialDst)
					cudaSafeCall(cudaMalloc(
					  &deviceDstStride, sizeof(int64_t) * LIBRAPID_MAX_DIMS));
				cudaSafeCall(cudaMalloc(&deviceSrcStride,
										sizeof(int64_t) * LIBRAPID_MAX_DIMS));

				cudaSafeCall(cudaMemcpy(deviceExtent,
										extent.raw(),
										sizeof(int64_t) * LIBRAPID_MAX_DIMS,
										cudaMemcpyHostToDevice));
				if (!trivialDst)
					cudaSafeCall(cudaMemcpy(deviceDstStride,
											dstStride.raw(),
											sizeof(int64_t) * LIBRAPID_MAX_DIMS,
											cudaMemcpyHostToDevice));
				cudaSafeCall(cudaMemcpy(deviceSrcStride,
										srcStride.raw(),
										sizeof(int64_t) * LIBRAPID_MAX_DIMS,
										cudaMemcpyHostToDevice));
#	endif // LIBRAPID_CUDA_STREAM

				std::string kernel = "unaryKernelComplex\n";
				kernel +=
				  "#include <stdint.h>\n"
				  "#include <type_traits>\n"
				  "#include <" CUDA_INCLUDE_DIRS
				  "/curand_kernel.h>\n"
				  "#include <" CUDA_INCLUDE_DIRS "/curand.h>\n\n";

				kernel += "const int64_t LIBRAPID_MAX_DIMS = " +
						  std::to_string(LIBRAPID_MAX_DIMS) + ";";
				kernel += R"V0G0N(
					__device__
					inline int64_t indexToIndex(uint64_t index,
											   const int64_t *__restrict shape,
											   const int64_t *__restrict strides,
											   uint64_t dims)
					{
						int64_t products[LIBRAPID_MAX_DIMS];
						int64_t prod = 1;
						for (int64_t i = dims - 1; i >= 0; --i)
						{
							products[i] = prod;
							prod *= shape[i];
						}
						int64_t tmp = index;
						int64_t res = 0;
						int64_t tmp2;
						for (int64_t i = 0; i < dims; ++i)
						{
							tmp2 = tmp / products[i];
							res = res + strides[i] * tmp2;
							tmp = tmp - tmp2 * products[i];
						}
						return res;
					}
					)V0G0N";

				kernel += complexHpp;

				kernel += R"V0G0N(
					template<typename A>
					__device__
					inline auto )V0G0N";

				kernel += op.name;

				kernel +=
				  R"V0G0N((A &a, int64_t indexA, curandState_t *_curandState)
					{
					)V0G0N";

				kernel += op.kernel;

				kernel += R"V0G0N(}
					template<typename T_DST, typename T_SRC>
					__global__
						void unaryFuncComplex(T_DST * __restrict dstData,
											  const T_SRC * __restrict srcData, int64_t size,
											  const int64_t extent[LIBRAPID_MAX_DIMS],
											  const int64_t dstStride[LIBRAPID_MAX_DIMS],
											  const int64_t srcStride[LIBRAPID_MAX_DIMS],
											  int64_t dims,
											  curandState_t *_curandStates = nullptr)
					{
						uint64_t kernelIndex = blockDim.x * blockIdx.x
											   + threadIdx.x;
						uint64_t srcIndex = indexToIndex(kernelIndex, extent, srcStride, dims);
					)V0G0N";

				if (trivialDst)
					kernel += "uint64_t dstIndex = kernelIndex;";
				else
					kernel +=
					  "uint64_t dstIndex = indexToIndex(kernelIndex, extent, "
					  "dstStride, dims);";

				kernel += R"V0G0N(
						if (kernelIndex < size) {
					)V0G0N";

				kernel +=
				  "dstData[dstIndex] = " + op.name +
				  "(srcData[srcIndex], srcIndex, _curandStates + kernelIndex);";
				kernel += "\n}\n}";

				static jitify::JitCache kernelCache;
				jitify::Program program =
				  kernelCache.program(kernel, cudaHeaders, cudaParams);

				unsigned int threadsPerBlock, blocksPerGrid;

				// Use 1 to 512 threads per block
				if (elems < 512) {
					threadsPerBlock = (unsigned int)elems;
					blocksPerGrid	= 1;
				} else {
					threadsPerBlock = 512;
					blocksPerGrid =
					  ceil(double(elems) / double(threadsPerBlock));
				}

				dim3 grid(blocksPerGrid);
				dim3 block(threadsPerBlock);

				static curandState_t *curandGenerators = nullptr;
				static int64_t numGenerators		   = 0;
				if ((op.name == "fillRandom" ||
					 op.name == "fillRandomComplex") &&
					(numGenerators < threadsPerBlock * blocksPerGrid)) {
					// Free curand generators and reallocate with a larger size
					// Allocate the nearest power of two generators
					cudaSafeCall(cudaFree(curandGenerators));

					numGenerators = 1;
					while (numGenerators < threadsPerBlock * blocksPerGrid)
						numGenerators <<= 1;

#	ifdef LIBRAPID_CUDA_STREAM
					cudaSafeCall(
					  cudaMallocAsync(&curandGenerators,
									  sizeof(curandState_t) * numGenerators,
									  cudaStream));
#	else
					cudaSafeCall(
					  cudaMallocAsync(&curandGenerators,
									  sizeof(curandState_t) * numGenerators));
#	endif // LIBRAPID_CUDA_STREAM

					std::string randomKernel =
					  std::string("randomKernel\n") + "#include <stdint.h>\n" +
					  "#include <type_traits>\n" + "#include <" +
					  CUDA_INCLUDE_DIRS + "/curand_kernel.h>\n" + "#include <" +
					  CUDA_INCLUDE_DIRS + "/curand.h>\n\n";

					randomKernel += R"V0G0N(
						__global__
						void populateGenerators(curandState_t *states, int64_t numStates) {
							int64_t kernelIndex = blockDim.x * blockIdx.x + threadIdx.x;
							curand_init()V0G0N";
					randomKernel +=
					  std::to_string(getSeed(op)) +
					  ", kernelIndex, 0, &states[kernelIndex]);}\n";

					static jitify::JitCache randomCache;
					jitify::Program randomProgram = kernelCache.program(
					  randomKernel, cudaHeaders, cudaParams);

#	ifdef LIBRAPID_CUDA_STREAM
					jitifyCall(randomProgram.kernel("populateGenerators")
								 .instantiate()
								 .configure(grid, block, 0, cudaStream)
								 .launch(curandGenerators, numGenerators));
#	else
					jitifyCall(program.kernel("populateGenerators")
								 .instantiate()
								 .configure(grid, block)
								 .launch(curandGenerators, numGenerators));
#	endif // LIBRAPID_CUDA_STREAM
				}

				std::visit(
				  [&](auto *dstData, auto *srcData) {
					  using T_DST =
						typename std::remove_pointer<decltype(dstData)>::type;
					  using T_SRC =
						typename std::remove_pointer<decltype(srcData)>::type;

#	ifdef LIBRAPID_CUDA_STREAM
					  jitifyCall(program.kernel("unaryFuncComplex")
								   .instantiate(Type<T_DST>(), Type<T_SRC>())
								   .configure(grid, block, 0, cudaStream)
								   .launch(dstData,
										   srcData,
										   elems,
										   deviceExtent,
										   deviceDstStride,
										   deviceSrcStride,
										   dims,
										   curandGenerators));
#	else
					  jitifyCall(program.kernel("unaryFuncComplex")
								   .instantiate(Type<T_DST>(), Type<T_SRC>())
								   .configure(grid, block)
								   .launch(dstData,
										   srcData,
										   elems,
										   deviceExtent,
										   deviceDstStride,
										   deviceSrcStride,
										   dims,
										   curandGenerators));
#	endif // LIBRAPID_CUDA_STREAM
				  },
				  dst.data,
				  src.data);

#	ifdef LIBRAPID_CUDA_STREAM
				cudaSafeCall(cudaFreeAsync(deviceExtent, cudaStream));
				if (!trivialDst)
					cudaSafeCall(cudaFreeAsync(deviceDstStride, cudaStream));
				cudaSafeCall(cudaFreeAsync(deviceSrcStride, cudaStream));
#	else
				cudaSafeCall(cudaDeviceSynchronize());
				cudaSafeCall(cudaFree(deviceExtent));
				if (!trivialDst) cudaSafeCall(cudaFree(deviceDstStride));
				cudaSafeCall(cudaFree(deviceSrcStride));
#	endif // LIBRAPID_CUDA_STREAM
			}
#else
			else {
				throw std::runtime_error("CUDA support was not enabled");
			}
#endif // LIBRAPID_HAS_CUDA
		}
	}

	/**
	 * \rst
	 *
	 * Perform a binary operation on trivially strided, contiguous memory.
	 *
	 * An operation struct can be passed to the ``op`` parameter of this
	 * function. It must contain a ``func``, being a lambda taking two
	 * templated values which returns a single ``auto`` value.
	 *
	 * If CUDA support is enabled, the operation struct must also contain a
	 * small kernel operation stored as a string. An example kernel is
	 * ``"c = a + b"`` -- "c", "b" and "a" are scalar values corresponding to
	 *the provided input data. This example kernel will result in the the data
	 *being added together, and the result being stored in ``c``
	 *
	 * If using scalar values as inputs, set the corresponding ``?IsScalar``
	 * parameter to true. Only one scalar input is allowed.
	 *
	 * Parameters
	 * ----------
	 * locnA: Accelerator
	 *		The location of ``a``
	 * locnB: Accelerator
	 *		The location of ``b``
	 * locnC: Accelerator
	 *		The location of the result data ``c``
	 * a: ``A *``
	 *		Trivial array
	 * b: ``B *``
	 *		Trivial array
	 * c: ``C *``
	 *		Trivial array
	 * aIsScalar: Boolean
	 *		If true, the input data ``a`` will be treated as a scalar
	 * bIsScalar: Boolean
	 *		If true, the input data ``b`` will be treated as a scalar
	 * size: integer
	 *		The number of *elements* in ``a`` (and therefore ``b``)
	 * op: Operation struct (containing "name", "kernel" and "func")
	 *
	 * \endrst
	 */
	template<typename FUNC>
	inline void multiarrayBinaryOpTrivial(RawArray &dst, const RawArray &srcA,
										  const RawArray &srcB,
										  bool srcAIsScalar, bool srcBIsScalar,
										  int64_t elems, const FUNC &op,
										  bool permitVectorize = true) {
		if (dst.location != srcA.location || dst.location != srcB.location) {
			// Locations are different, so make A and B have the same
			// accelerator as the result array (C)

			// Copy A and B to be on the same accelerator as the destination
			RawArray tempSrcA = {(int64_t *)nullptr, srcA.dtype, dst.location};
			RawArray tempSrcB = {(int64_t *)nullptr, srcB.dtype, dst.location};

			// Allocate memory for temporary sources
			// rawArrayMalloc(tempSrcA, srcAIsScalar ? 1 : elems);
			// rawArrayMalloc(tempSrcB, srcBIsScalar ? 1 : elems);

			// Copy the data from the original sources to the adjusted sources
			int freeSrcA =
			  makeSameAccelerator(tempSrcA, srcA, srcAIsScalar ? 1 : elems);
			int freeSrcB =
			  makeSameAccelerator(tempSrcB, srcB, srcBIsScalar ? 1 : elems);

			// Apply the operation again, using the new sources
			multiarrayBinaryOpTrivial(
			  dst, tempSrcA, tempSrcB, srcAIsScalar, srcBIsScalar, elems, op);

			// Free the allocated data -- may not be freed if a pointer was
			// simply copied. See the different free modes further up in this
			// file
			freeWithMode(tempSrcA, freeSrcA);
			freeWithMode(tempSrcB, freeSrcB);
		} else {
			// Locations are the same, so apply a single, unified operation
			if (dst.location == Accelerator::CPU) {
				std::visit(
				  [&](auto *dstData, auto *srcDataA, auto *srcDataB) {
					  auto tempOp	 = op;
					  auto tempElems = elems;

					  // Typenames which are useful for casting and checking
					  using C =
						typename std::remove_pointer<decltype(dstData)>::type;
					  using A =
						typename std::remove_pointer<decltype(srcDataA)>::type;
					  using B =
						typename std::remove_pointer<decltype(srcDataB)>::type;

					  if (srcAIsScalar) {
						  // Use *a rather than a[i]
						  if (elems < 2500) {
							  for (int64_t i = 0; i < elems; ++i)
								  dstData[i] = static_cast<C>(
									tempOp(*srcDataA, srcDataB[i], 0, i));
						  } else {
#pragma omp parallel for shared(                                               \
  dstData, srcDataA, srcDataB, tempElems, tempOp) default(none)
							  for (int64_t i = 0; i < tempElems; ++i) {
								  dstData[i] = static_cast<C>(
									tempOp(*srcDataA, srcDataB[i], 0, i));
							  }
						  }
					  } else if (srcBIsScalar) {
						  // Use *b rather than b[i]
						  if (elems < 2500) {
							  for (int64_t i = 0; i < tempElems; ++i)
								  dstData[i] = static_cast<C>(
									tempOp(srcDataA[i], *srcDataB, i, 0));
						  } else {
#pragma omp parallel for shared(                                               \
  dstData, srcDataA, srcDataB, tempElems, tempOp) default(none)
							  for (int64_t i = 0; i < tempElems; ++i) {
								  dstData[i] = static_cast<C>(
									tempOp(srcDataA[i], *srcDataB, i, 0));
							  }
						  }
					  } else {
						  // Use a[i] and b[i]
						  if (permitVectorize && std::is_same_v<A, double> &&
							  std::is_same_v<B, double> &&
							  std::is_same_v<C, double>) {
							  vcl::Vec8d a, b;
							  int64_t i	   = 0;
							  auto tmpSrcA = (double *__restrict)srcDataA;
							  auto tmpSrcB = (double *__restrict)srcDataB;
							  auto tmpDst  = (double *__restrict)dstData;

							  if (elems < 2500) {
								  for (i = 0; i < tempElems - 7; i += 8) {
									  a.load(tmpSrcA + i);
									  b.load(tmpSrcB + i);
									  vcl::Vec8d c = tempOp(a, b, i, i);
									  c.store(tmpDst + i);
								  }
							  } else {
#pragma omp parallel for shared(                                               \
  tmpDst, tmpSrcA, tmpSrcB, tempElems, tempOp, i) private(a, b) default(none)
								  for (i = 0; i < tempElems - 7; i += 8) {
									  a.load(tmpSrcA + i);
									  b.load(tmpSrcB + i);
									  vcl::Vec8d c = tempOp(a, b, i, i);
									  c.store(tmpDst + i);
								  }
							  }

							  int64_t diff = tempElems - i;
							  if (diff > 0) {
								  a.load_partial((int)diff, tmpSrcA + i);
								  b.load_partial((int)diff, tmpSrcB + i);
								  vcl::Vec8d c = tempOp(a, b, i, i);
								  c.store_partial((int)diff, tmpDst + i);
							  }
						  } else if (permitVectorize &&
									 std::is_same_v<A, float> &&
									 std::is_same_v<B, float> &&
									 std::is_same_v<C, float>) {
							  vcl::Vec16f a, b;
							  auto tmpSrcA = (float *__restrict)srcDataA;
							  auto tmpSrcB = (float *__restrict)srcDataB;
							  auto tmpDst  = (float *__restrict)dstData;

							  if (elems < 2500) {
								  for (int64_t i = 0; i < tempElems - 15;
									   i += 16) {
									  a.load(tmpSrcA + i);
									  b.load(tmpSrcB + i);
									  vcl::Vec16f c = tempOp(a, b, i, i);
									  c.store(tmpDst + i);
								  }
							  } else {
#pragma omp parallel for shared(                                               \
  tmpDst, tmpSrcA, tmpSrcB, tempElems, tempOp) private(a, b) default(none)
								  for (int64_t i = 0; i < tempElems - 15;
									   i += 16) {
									  a.load(tmpSrcA + i);
									  b.load(tmpSrcB + i);
									  vcl::Vec16f c = tempOp(a, b, i, i);
									  c.store(tmpDst + i);
								  }
							  }

							  if (elems & 15) {
								  a.load(tmpSrcA + elems - 16);
								  b.load(tmpSrcB + elems - 16);
								  vcl::Vec16f c =
									tempOp(a, b, elems - 16, elems - 16);
								  c.store(tmpDst + elems - 16);
							  }
						  } else if (elems < 2500) {
							  for (int64_t i = 0; i < tempElems; ++i) {
								  dstData[i] = static_cast<C>(
									tempOp(srcDataA[i], srcDataB[i], i, i));
							  }
						  } else {
#pragma omp parallel for shared(                                               \
  dstData, srcDataA, srcDataB, tempElems, tempOp) default(none)
							  for (int64_t i = 0; i < tempElems; ++i) {
								  dstData[i] = static_cast<C>(
									tempOp(srcDataA[i], srcDataB[i], i, i));
							  }
						  }
					  }
				  },
				  dst.data,
				  srcA.data,
				  srcB.data);
			}
#ifdef LIBRAPID_HAS_CUDA
			else {
				using jitify::reflection::Type;

				std::string kernel = "binaryKernelTrivial\n";
				kernel += "__constant__ int LIBRAPID_MAX_DIMS = " +
						  std::to_string(LIBRAPID_MAX_DIMS) + ";\n";
				kernel +=
				  "#include <stdint.h>\n"
				  "#include <type_traits>\n"
				  "#include <" CUDA_INCLUDE_DIRS
				  "/curand_kernel.h>\n"
				  "#include <" CUDA_INCLUDE_DIRS "/curand.h>\n\n";

				kernel += complexHpp;

				// Insert the user-defined kernel into the main GPU kernel
				kernel += R"V0G0N(
					template<typename A, typename B>
					__device__
					inline auto )V0G0N";

				kernel += op.name;

				kernel += R"V0G0N((A &a, B &b, int64_t indexA, int64_t indexB)
					{
					)V0G0N";

				kernel += op.kernel;

				kernel += R"V0G0N(}
                    template<typename T, typename std::enable_if<std::is_floating_point<T>::value, int>::type = 0>
                    __global__
                    inline T random(T lower, T upper, uint64_t seed = -1) {
                        // Random floating point value in range [lower, upper)
                        static std::uniform_real_distribution<T> distribution(0., 1.);
                        static std::mt19937 generator(seed == (uint64_t) -1 ? (unsigned int) (seconds() * 10) : seed);
                        return lower + (upper - lower) * distribution(generator);
                    }
                    template<typename T, typename std::enable_if<!std::is_floating_point<T>::value, int>::type = 0>
                    __global__
                    inline T random(T lower, T upper, uint64_t seed = -1) {
                        // Random integral value in range [lower, upper]
                        return (T) random((double) (lower - (lower < 0 ? 1 : 0)), (double) upper + 1, seed);
                    }
					template<typename T_DST, typename T_SRCA, typename T_SRCB>
					__global__
					void binaryFuncTrivial(T_DST *__restrict dstData,
										   const T_SRCA *__restrict srcA,
										   const T_SRCB *__restrict srcB,
										   int64_t size)
					{
						const int64_t kernelIndex = blockDim.x * blockIdx.x
												   + threadIdx.x;
						if (kernelIndex < size) {
					)V0G0N";

				if (srcAIsScalar)
					kernel += "dstData[kernelIndex] = " + op.name +
							  "(*srcA, srcB[kernelIndex], 0, kernelIndex);";
				else if (srcBIsScalar)
					kernel += "dstData[kernelIndex] = " + op.name +
							  "(srcA[kernelIndex], *srcB, kernelIndex, 0);";
				else
					kernel += "dstData[kernelIndex] = " + op.name +
							  "(srcA[kernelIndex], srcB[kernelIndex], "
							  "kernelIndex, kernelIndex);";
				kernel += "\n}\n}";

				// fmt::print("Information:\n{}\n", kernel);

				static const std::vector<std::string> params = {
				  "--disable-warnings",
				  "-std=c++17",
				  std::string("-I \"") + CUDA_INCLUDE_DIRS + "\""};

				static jitify::JitCache kernelCache;
				jitify::Program program =
				  kernelCache.program(kernel, cudaHeaders, params);

				int64_t threadsPerBlock, blocksPerGrid;

				// Use 1 to 512 threads per block
				if (elems < 512) {
					threadsPerBlock = elems;
					blocksPerGrid	= 1;
				} else {
					threadsPerBlock = 512;
					blocksPerGrid =
					  ceil(double(elems) / double(threadsPerBlock));
				}

				dim3 grid(blocksPerGrid);
				dim3 block(threadsPerBlock);

				std::visit(
				  [&](auto *__restrict dstData,
					  auto *__restrict srcDataA,
					  auto *__restrict srcDataB) {
					  using T_DST =
						typename std::remove_pointer_t<decltype(dstData)>;
					  using T_SRCA =
						typename std::remove_pointer_t<decltype(srcDataA)>;
					  using T_SRCB =
						typename std::remove_pointer_t<decltype(srcDataB)>;

					// auto thing = (T_DST) 0;
					// std::cout << "THING: " << thing << "\n";

#	ifdef LIBRAPID_CUDA_STREAM
					  jitifyCall(program.kernel("binaryFuncTrivial")
								   .instantiate(Type<T_DST>(),
												Type<T_SRCA>(),
												Type<T_SRCB>())
								   .configure(grid, block, 0, cudaStream)
								   .launch(dstData, srcDataA, srcDataB, elems));
#	else
					  jitifyCall(program.kernel("binaryFuncTrivial")
								   .instantiate(Type<T_DST>(),
												Type<T_SRCA>(),
												Type<T_SRCB>())
								   .configure(grid, block)
								   .launch(dstData, srcDataA, srcDataB, elems));
#	endif // LIBRAPID_CUDA_STREAM
				  },
				  dst.data,
				  srcA.data,
				  srcB.data);
			}
#endif // LIBRAPID_HAS_CUDA
		}
	}

	template<typename FUNC>
	inline void multiarrayBinaryOpComplex(
	  RawArray &dst, const RawArray &srcA, const RawArray &srcB,
	  bool srcAIsScalar, bool srcBIsScalar, const int64_t elems,
	  const Extent &extent, const Stride &strideDst, const Stride &strideSrcA,
	  const Stride &strideSrcB, const FUNC &op) {
		if (dst.location != srcA.location || dst.location != srcB.location) {
			// Locations are different, so make A and B have the same
			// accelerator as the result array (C)

			// Copy A to be on the same accelerator as B
			RawArray tempSrcA = {(int64_t *)nullptr, dst.dtype, dst.location};
			RawArray tempSrcB = {(int64_t *)nullptr, dst.dtype, dst.location};

			// Allocate memory for temporary sources
			rawArrayMalloc(tempSrcA, elems);
			rawArrayMalloc(tempSrcB, elems);

			// Copy the data from the original sources to the adjusted sources
			int freeSrcA =
			  makeSameAccelerator(tempSrcA, srcA, srcAIsScalar ? 1 : elems);
			int freeSrcB =
			  makeSameAccelerator(tempSrcB, srcB, srcBIsScalar ? 1 : elems);

			// Apply the operation again, using the new sources
			multiarrayBinaryOpComplex(dst,
									  tempSrcA,
									  tempSrcB,
									  srcAIsScalar,
									  srcBIsScalar,
									  elems,
									  extent,
									  strideDst,
									  strideSrcA,
									  strideSrcB,
									  op);

			// Free the allocated data -- may not be freed if a pointer was
			// simply copied. See the different free modes further up in this
			// file
			freeWithMode(tempSrcA, freeSrcA);
			freeWithMode(tempSrcB, freeSrcB);
		} else {
			if (dst.location == Accelerator::CPU) {
				// Iterate over the array using its stride and extent
				int64_t coord[LIBRAPID_MAX_DIMS] {};

				// Counters
				int64_t idim = 0;
				int64_t ndim = extent.ndim();

				// Create pointers here so repeated function calls
				// are not needed
				static int64_t rawExtent[LIBRAPID_MAX_DIMS];
				static int64_t _strideDst[LIBRAPID_MAX_DIMS];
				static int64_t _strideSrcA[LIBRAPID_MAX_DIMS];
				static int64_t _strideSrcB[LIBRAPID_MAX_DIMS];

				for (int64_t i = 0; i < ndim; ++i) {
					rawExtent[ndim - i - 1]	  = extent[i];
					_strideDst[ndim - i - 1]  = strideDst[i];
					_strideSrcA[ndim - i - 1] = strideSrcA[i];
					_strideSrcB[ndim - i - 1] = strideSrcB[i];
				}

				std::visit(
				  [&](auto *dstData, auto *srcA, auto *srcB) {
					  using C =
						typename std::remove_pointer<decltype(dstData)>::type;
					  using A =
						typename std::remove_pointer<decltype(srcA)>::type;
					  using B =
						typename std::remove_pointer<decltype(srcB)>::type;

					  int64_t dstIndex = 0, srcAIndex = 0, srcBIndex = 0;
					  if (srcAIsScalar) {
						  // Use *a rather than a[i]
						  do {
							  // *dstData = (C) op(*srcA, *srcB);
							  dstData[dstIndex] =
								(C)op(*srcA, srcB[srcBIndex], 0, srcBIndex);

							  for (idim = 0; idim < ndim; ++idim) {
								  if (++coord[idim] == rawExtent[idim]) {
									  coord[idim] = 0;
									  // srcB = srcB - (rawExtent[idim] - 1) *
									  // _strideSrcB[idim]; dstData = dstData -
									  // (rawExtent[idim] - 1) *
									  // _strideDst[idim];
									  srcBIndex =
										srcBIndex - (rawExtent[idim] - 1) *
													  _strideSrcB[idim];
									  dstIndex =
										dstIndex - (rawExtent[idim] - 1) *
													 _strideDst[idim];
								  } else {
									  // srcB = srcB + _strideSrcB[idim];
									  // dstData = dstData + _strideDst[idim];
									  srcBIndex = srcBIndex + _strideSrcB[idim];
									  dstIndex	= dstIndex + _strideDst[idim];
									  break;
								  }
							  }
						  } while (idim < ndim);
					  } else if (srcBIsScalar) {
						  // Use *b rather than b[i]
						  do {
							  // *dstData = (C) op(*srcA, *srcB);
							  dstData[dstIndex] =
								op(srcA[srcAIndex], *srcB, srcAIndex, 0);

							  for (idim = 0; idim < ndim; ++idim) {
								  if (++coord[idim] == rawExtent[idim]) {
									  coord[idim] = 0;
									  // srcA = srcA - (rawExtent[idim] - 1) *
									  // _strideSrcA[idim]; dstData = dstData -
									  // (rawExtent[idim] - 1) *
									  // _strideDst[idim];
									  srcAIndex =
										srcAIndex - (rawExtent[idim] - 1) *
													  _strideSrcA[idim];
									  dstIndex =
										dstIndex - (rawExtent[idim] - 1) *
													 _strideDst[idim];
								  } else {
									  // srcA = srcA + _strideSrcA[idim];
									  // dstData = dstData + _strideDst[idim];
									  srcAIndex = srcAIndex + _strideSrcA[idim];
									  dstIndex	= dstIndex + _strideDst[idim];
									  break;
								  }
							  }
						  } while (idim < ndim);
					  } else {
						  // Use a[i] and b[i]
						  do {
							  // *dstData = (C) op(*srcA, *srcB);
							  dstData[dstIndex] = (C)op(srcA[srcAIndex],
														srcB[srcBIndex],
														srcAIndex,
														srcBIndex);

							  for (idim = 0; idim < ndim; ++idim) {
								  if (++coord[idim] == rawExtent[idim]) {
									  coord[idim] = 0;
									  // dstData = dstData - (rawExtent[idim] -
									  // 1) * _strideDst[idim]; srcA = srcA -
									  // (rawExtent[idim] - 1) *
									  // _strideSrcA[idim]; srcB = srcB -
									  // (rawExtent[idim] - 1) *
									  // _strideSrcB[idim];
									  dstIndex =
										dstIndex - (rawExtent[idim] - 1) *
													 _strideDst[idim];
									  srcAIndex =
										srcAIndex - (rawExtent[idim] - 1) *
													  _strideSrcA[idim];
									  srcB = srcB - (rawExtent[idim] - 1) *
													  _strideSrcB[idim];
								  } else {
									  // dstData = dstData + _strideDst[idim];
									  // srcA = srcA + _strideSrcA[idim];
									  // srcB = srcB + _strideSrcB[idim];
									  dstIndex	= dstIndex + _strideDst[idim];
									  srcAIndex = srcAIndex + _strideSrcA[idim];
									  srcBIndex = srcBIndex + _strideSrcB[idim];
									  break;
								  }
							  }
						  } while (idim < ndim);
					  }
				  },
				  dst.data,
				  srcA.data,
				  srcB.data);
			}
#ifdef LIBRAPID_HAS_CUDA
			else {
				using jitify::reflection::Type;

				// Buffers for extent and strides
				int64_t dims = extent.ndim();
				int64_t *deviceExtent, *deviceStrideSrcA, *deviceStrideSrcB,
				  *deviceStrideDst;

#	ifdef LIBRAPID_CUDA_STREAM
				cudaSafeCall(
				  cudaMallocAsync(&deviceExtent,
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaStream));
				cudaSafeCall(
				  cudaMallocAsync(&deviceStrideSrcA,
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaStream));
				cudaSafeCall(
				  cudaMallocAsync(&deviceStrideSrcB,
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaStream));
				cudaSafeCall(
				  cudaMallocAsync(&deviceStrideDst,
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaStream));

				cudaSafeCall(
				  cudaMemcpyAsync(deviceExtent,
								  extent.raw(),
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaMemcpyHostToDevice,
								  cudaStream));
				cudaSafeCall(
				  cudaMemcpyAsync(deviceStrideSrcA,
								  strideSrcA.raw(),
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaMemcpyHostToDevice,
								  cudaStream));
				cudaSafeCall(
				  cudaMemcpyAsync(deviceStrideSrcB,
								  strideSrcB.raw(),
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaMemcpyHostToDevice,
								  cudaStream));
				cudaSafeCall(
				  cudaMemcpyAsync(deviceStrideDst,
								  strideDst.raw(),
								  sizeof(int64_t) * LIBRAPID_MAX_DIMS,
								  cudaMemcpyHostToDevice,
								  cudaStream));
#	else
				cudaSafeCall(cudaMalloc(&deviceExtent,
										sizeof(int64_t) * LIBRAPID_MAX_DIMS));
				cudaSafeCall(cudaMalloc(&deviceStrideSrcA,
										sizeof(int64_t) * LIBRAPID_MAX_DIMS));
				cudaSafeCall(cudaMalloc(&deviceStrideSrcB,
										sizeof(int64_t) * LIBRAPID_MAX_DIMS));
				cudaSafeCall(cudaMalloc(&deviceStrideDst,
										sizeof(int64_t) * LIBRAPID_MAX_DIMS));

				cudaSafeCall(cudaMemcpy(deviceExtent,
										extent.raw(),
										sizeof(int64_t) * LIBRAPID_MAX_DIMS,
										cudaMemcpyHostToDevice));
				cudaSafeCall(cudaMemcpy(deviceStrideSrcA,
										_strideSrcA.raw(),
										sizeof(int64_t) * LIBRAPID_MAX_DIMS,
										cudaMemcpyHostToDevice));
				cudaSafeCall(cudaMemcpy(deviceStrideSrcB,
										strideSrcB.raw(),
										sizeof(int64_t) * LIBRAPID_MAX_DIMS,
										cudaMemcpyHostToDevice));
				cudaSafeCall(cudaMemcpy(deviceStrideDst,
										strideDst.raw(),
										sizeof(int64_t) * LIBRAPID_MAX_DIMS,
										cudaMemcpyHostToDevice));
#	endif // LIBRAPID_CUDA_STREAM

				std::string kernel = "binaryKernelComplex\n";
				kernel += "const int64_t LIBRAPID_MAX_DIMS = " +
						  std::to_string(LIBRAPID_MAX_DIMS) + ";\n";

				kernel +=
				  "#include <stdint.h>\n"
				  "#include <type_traits>\n"
				  "#include <" CUDA_INCLUDE_DIRS
				  "/curand_kernel.h>\n"
				  "#include <" CUDA_INCLUDE_DIRS "/curand.h>\n\n";

				kernel += R"V0G0N(
                    template<typename T, typename std::enable_if<std::is_floating_point<T>::value, int>::type = 0>
                    __global__
                    inline T random(T lower, T upper, uint64_t seed = -1) {
                        // Random floating point value in range [lower, upper)
                        static std::uniform_real_distribution<T> distribution(0., 1.);
                        static std::mt19937 generator(seed == (uint64_t) -1 ? (unsigned int) (seconds() * 10) : seed);
                        return lower + (upper - lower) * distribution(generator);
                    }
                    template<typename T, typename std::enable_if<!std::is_floating_point<T>::value, int>::type = 0>
                    __global__
                    inline T random(T lower, T upper, uint64_t seed = -1) {
                        // Random integral value in range [lower, upper]
                        return (T) random((double) (lower - (lower < 0 ? 1 : 0)), (double) upper + 1, seed);
                    }
					__device__
					inline int64_t indexToIndex(uint64_t index,
											   const int64_t *__restrict shape,
											   const int64_t *__restrict strides,
											   uint64_t dims)
					{
						int64_t products[LIBRAPID_MAX_DIMS];
						int64_t prod = 1;
						for (int64_t i = dims - 1; i >= 0; --i)
						{
							products[i] = prod;
							prod *= shape[i];
						}
						int64_t tmp = index;
						int64_t res = 0;
						int64_t tmp2;
						for (int64_t i = 0; i < dims; ++i)
						{
							tmp2 = tmp / products[i];
							res = res + strides[i] * tmp2;
							tmp = tmp - tmp2 * products[i];
						}
						return res;
					}
					)V0G0N";

				kernel += complexHpp;

				kernel += R"V0G0N(
					template<typename A, typename B>
					__device__
					inline auto )V0G0N";

				kernel += op.name;

				kernel += R"V0G0N((A &a, B &b, int64_t indexA, int64_t indexB)
					{
					)V0G0N";

				kernel += op.kernel;

				kernel += R"V0G0N(}
					template<typename T_DST, typename T_SRCA, typename T_SRCB>
					__global__
					void binaryFuncComplex(T_DST *__restrict dstData,
										   const T_SRCA *__restrict srcA,
										   const T_SRCB *__restrict srcB,
										   int64_t size,
										   const int64_t extent[LIBRAPID_MAX_DIMS],
										   const int64_t strideC[LIBRAPID_MAX_DIMS],
										   const int64_t strideA[LIBRAPID_MAX_DIMS],
										   const int64_t strideB[LIBRAPID_MAX_DIMS],
										   const int64_t dims)
					{
						int64_t kernelIndex = blockDim.x * blockIdx.x
														+ threadIdx.x;
						int64_t indexDst = indexToIndex(kernelIndex, extent, strideC, dims);
						int64_t indexSrcA = indexToIndex(kernelIndex, extent, strideA, dims);
						int64_t indexSrcB = indexToIndex(kernelIndex, extent, strideB, dims);
						if (kernelIndex < size) {
					)V0G0N";

				if (srcAIsScalar)
					kernel += "dstData[indexDst] = " + op.name +
							  "(*srcA, srcB[indexSrcB], 0, indexSrcB);";
				else if (srcBIsScalar)
					kernel += "dstData[indexDst] = " + op.name +
							  "(srcA[indexSrcA], *srcB, indexSrcA, 0);";
				else
					kernel += "dstData[indexDst] = " + op.name +
							  "(srcA[indexSrcA], srcB[indexSrcB], indexSrcA, "
							  "indexSrcB);";

				kernel += "\n}\n}";

				static jitify::JitCache kernelCache;
				jitify::Program program =
				  kernelCache.program(kernel, cudaHeaders, cudaParams);

				unsigned int threadsPerBlock, blocksPerGrid;

				// Use 1 to 512 threads per block
				if (elems < 512) {
					threadsPerBlock = (unsigned int)elems;
					blocksPerGrid	= 1;
				} else {
					threadsPerBlock = 512;
					blocksPerGrid =
					  ceil(double(elems) / double(threadsPerBlock));
				}

				dim3 grid(blocksPerGrid);
				dim3 block(threadsPerBlock);

				std::visit(
				  [&](auto *dstData, auto *srcDataA, auto *srcDataB) {
					  using T_SRCA =
						typename std::remove_pointer<decltype(srcDataA)>::type;
					  using T_SRCB =
						typename std::remove_pointer<decltype(srcDataB)>::type;
					  using T_DST =
						typename std::remove_pointer<decltype(dstData)>::type;

#	ifdef LIBRAPID_CUDA_STREAM
					  jitifyCall(program.kernel("binaryFuncComplex")
								   .instantiate(Type<T_DST>(),
												Type<T_SRCA>(),
												Type<T_SRCB>())
								   .configure(grid, block, 0, cudaStream)
								   .launch(dstData,
										   srcDataA,
										   srcDataB,
										   elems,
										   deviceExtent,
										   deviceStrideDst,
										   deviceStrideSrcA,
										   deviceStrideSrcB,
										   dims));
#	else
					  jitifyCall(program.kernel("binaryFuncComplex")
								   .instantiate(Type<T_DST>(),
												Type<T_SRCA>(),
												Type<T_SRCB>())
								   .configure(grid, block)
								   .launch(dstData,
										   srcDataA,
										   srcDataB,
										   elems,
										   deviceExtent,
										   deviceStrideSrcA,
										   deviceStrideSrcB,
										   deviceStrideDst,
										   dims));
#	endif // LIBRAPID_CUDA_STREAM
				  },
				  dst.data,
				  srcA.data,
				  srcB.data);

#	ifdef LIBRAPID_CUDA_STREAM
				cudaSafeCall(cudaFreeAsync(deviceExtent, cudaStream));
				cudaSafeCall(cudaFreeAsync(deviceStrideSrcA, cudaStream));
				cudaSafeCall(cudaFreeAsync(deviceStrideSrcB, cudaStream));
				cudaSafeCall(cudaFreeAsync(deviceStrideDst, cudaStream));
#	else
				cudaSafeCall(cudaFree(deviceExtent));
				cudaSafeCall(cudaFree(deviceStrideSrcA));
				cudaSafeCall(cudaFree(deviceStrideSrcB));
				cudaSafeCall(cudaFree(deviceStrideDst));
#	endif // LIBRAPID_CUDA_STREAM
			}
#endif // LIBRAPID_HAS_CUDA
		}
	}

} // namespace librapid::imp

#endif // LIBRAPID_MUTLIARRAY_OPERATIONS