#ifndef STIM_CUDA_GRADIENT3_H
#define STIM_CUDA_GRADIENT3_H

#include <iostream>
#include <cuda.h>
#include <stim/cuda/cudatools.h>
#include <stim/cuda/cudatools/error.h>

template<typename T>
__global__ void gradient3(T* out, T* in, int x, int y, int z){

	//calculate x,y,z coordinates for this thread
	int xi = blockIdx.x * blockDim.x + threadIdx.x;
	//find the grid size along y
	int grid_y = y / blockDim.y;
	int blockidx_y = blockIdx.y % grid_y;
	int yi = blockidx_y * blockDim.y + threadIdx.y;
	int zi = blockIdx.y / grid_y;
	int i = zi * x * y + yi * x + xi;
	
	//return if the pixel is outside of the image
	if(xi >= x || yi >= y || zi>=z) return;

	//calculate indices for the forward difference
	int i_xp = zi * x * y + yi * x + (xi + 1);
	int i_yp = zi * x * y + (yi + 1) * x + xi;
	int i_zp = (zi + 1) * x * y + yi * x + xi;

	//calculate indices for the backward difference
	int i_xn = zi * x * y + yi * x + (xi - 1);
	int i_yn = zi * x * y + (yi - 1) * x + xi;
	int i_zn = (zi - 1) * x * y + yi * x + xi;

	//use forward differences if a coordinate is zero
	if(xi == 0)
		out[i * 3 + 0] = in[i_xp] - in[i];
	if(yi == 0)
		out[i * 3 + 1] = in[i_yp] - in[i];
	if (zi==0)
		out[i * 3 + 2] = in[i_zp] - in[i];

	//use backward differences if the coordinate is at the maximum edge
	if(xi == x-1)
		out[i * 3 + 0] = in[i] - in[i_xn];
	if(yi == y-1)
		out[i * 3 + 1] = in[i] - in[i_yn];
	if(zi == z-1)
		out[i * 3 + 2] = in[i] - in[i_zn];

	//otherwise use central differences
	if(xi > 0 && xi < x-1)
		out[i * 3 + 0] = (in[i_xp] - in[i_xn]) / 2;
	if(yi > 0 && yi < y-1)
		out[i * 3 + 1] = (in[i_yp] - in[i_yn]) / 2;
	if(zi > 0 && zi < z-1)
		out[i * 3 + 2] = (in[i_zp] - in[i_zn]) / 2;

}

template<typename T>

void gpu_gradient3(T* gpuGrad, T* gpuI, size_t x, size_t y, size_t z){

			
	unsigned int max_threads = stim::maxThreadsPerBlock();
	dim3 threads((unsigned int)sqrt (max_threads), (unsigned int)sqrt (max_threads));
	dim3 blocks((unsigned int)x / threads.x + 1, ((unsigned int)y / threads.y + 1) * (unsigned int)z);
	
	//call the GPU kernel to determine the gradient
	gradient3<T> <<< blocks, threads >>>(gpuGrad, gpuI, (int)x, (int)y, (int)z);

}

template<typename T>
void cpu_gradient3(T* out, T* in, size_t x, size_t y, size_t z){

	//get the number of pixels in the image
	size_t pixels = x * y * z;
	size_t bytes = pixels * sizeof(T);

	//allocate space on the GPU for the input image
	T* gpuIn;
	HANDLE_ERROR(cudaMalloc(&gpuIn, bytes));

	//copy the image data to the GPU
	HANDLE_ERROR(cudaMemcpy(gpuIn, in, bytes, cudaMemcpyHostToDevice));

	//allocate space on the GPU for the output gradient image
	T* gpuOut;
	cudaMalloc(&gpuOut, bytes * 3);		//the output image will have two channels (x, y)

	//call the GPU version of this function
	gpu_gradient3(gpuOut, gpuIn, x, y, z);	

	//copy the results to the CPU
	cudaMemcpy(out, gpuOut, bytes * 3, cudaMemcpyDeviceToHost);

	//free allocated memory
	cudaFree(gpuIn);
	cudaFree(gpuOut);
}



#endif