#ifndef STIM_CUDA_GRADIENT_H
#define STIM_CUDA_GRADIENT_H

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

namespace stim{
	namespace cuda{

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

			
			// calculate the 2D coordinates for this current thread.
			size_t xi = blockIdx.x * blockDim.x + threadIdx.x;
			size_t yi = blockIdx.y * blockDim.y + threadIdx.y;
			// convert 2D coordinates to 1D
			size_t i = yi * x + xi;

			//return if the pixel is outside of the image
			if(xi >= x || yi >= y) return;

			//calculate indices for the forward difference
			size_t i_xp = yi * x + (xi + 1);
			size_t i_yp = (yi + 1) * x + xi;

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

			//use forward differences if a coordinate is zero
			if(xi == 0)
				out[i * 2 + 0] = in[i_xp] - in[i];
			else if (xi == x - 1)
				out[i * 2 + 0] = in[i] - in[i_xn];
			else
				out[i * 2 + 0] = (in[i_xp] - in[i_xn]) / 2;

			if(yi == 0)
				out[i * 2 + 1] = in[i_yp] - in[i];
			else if(yi == y-1)
				out[i * 2 + 1] = in[i] - in[i_yn];
			else
				out[i * 2 + 1] = (in[i_yp] - in[i_yn]) / 2;

		}

		template<typename T>
		void gpu_gradient_2d(T* gpuGrad, T* gpuI, size_t x, size_t y){
			
			//get the maximum number of threads per block for the CUDA device
			unsigned int max_threads = stim::maxThreadsPerBlock();
			dim3 threads(max_threads, 1);
			dim3 blocks((unsigned int)(x/threads.x) + 1 , (unsigned int)y);
			

			//call the GPU kernel to determine the gradient
			gradient_2d<T> <<< blocks, threads >>>(gpuGrad, gpuI, x, y);

		}

		template<typename T>
		void cpu_gradient_2d(T* out, T* in, unsigned int x, unsigned int y){

			//get the number of pixels in the image
			unsigned int pixels = x * y;
			unsigned int 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 * 2);		//the output image will have two channels (x, y)

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

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

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

	}
}


#endif