#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, unsigned int x, unsigned int y){

			//calculate the 1D image index for this thread
			int i = blockIdx.x * blockDim.x + threadIdx.x;

			//convert this to 2D pixel coordinates
			int yi = i / x;
			int xi = i - (yi * x);

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

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

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

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

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

			//otherwise use central differences
			if(xi > 0 && xi < x-1)
				out[i * 2 + 0] = (in[i_xp] - in[i_xn]) / 2;

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

		}

		template<typename T>
		//void gpu_gradient_2d(T* gpuOut, T* gpuIn, unsigned int x, unsigned int y){
		void gpu_gradient_2d(T* gpuGrad, T* gpuI, unsigned int x, unsigned int y){

			//get the number of pixels in the image
			unsigned int pixels = x * y;
			
			//allocate space on the GPU for the input image
			//T* gpuI;
			//HANDLE_ERROR(cudaMalloc(&gpuI, bytes));

			//cudaMemcpy(gpuI, gpuI0, bytes, cudaMemcpyDeviceToDevice);


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

			//get the maximum number of threads per block for the CUDA device
			int threads = stim::maxThreadsPerBlock();

			//calculate the number of blocks
			int blocks = pixels / threads + (pixels%threads == 0 ? 0:1);

			//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);
		}

	}
}


#endif