#ifndef RTS_COLORMAP_H
#define RTS_COLORMAP_H

#include <string>
#include <qimage.h>
#include <qcolor.h>
#include "rts/cuda/error.h"


#define BREWER_CTRL_PTS 11

#ifdef __CUDACC__
texture<float4, cudaTextureType1D> cudaTexBrewer;
static cudaArray* gpuBrewer;
#endif



namespace rts{
	namespace colormap{

enum colormapType {cmBrewer, cmGrayscale};

static void buffer2image(unsigned char* buffer, std::string filename, unsigned int x_size, unsigned int y_size)
{
	//create an image object
	QImage image(x_size, y_size, QImage::Format_RGB32);

	int i;
	unsigned char r, g, b;
	unsigned int x, y;
	for(y=0; y<y_size; y++)
		for(x=0; x<x_size; x++)
		{
			//calculate the 1D index
			i = y * x_size + x;

			r = buffer[i * 3 + 0];
			g = buffer[i * 3 + 1];
			b = buffer[i * 3 + 2];

			//set the image pixel
			QColor color(r, g, b);
			image.setPixel(x, y, color.rgb());
		}

	image.save(filename.c_str());
}

#ifdef __CUDACC__
static void initBrewer()
{
	//initialize the Brewer colormap

	//allocate CPU space
	float4 cpuColorMap[BREWER_CTRL_PTS];

	//define control rtsPoints
	cpuColorMap[0] = make_float4(0.192157f, 0.211765f, 0.584314f, 1.0f);
	cpuColorMap[1] = make_float4(0.270588f, 0.458824f, 0.705882f, 1.0f);
	cpuColorMap[2] = make_float4(0.454902f, 0.678431f, 0.819608f, 1.0f);
	cpuColorMap[3] = make_float4(0.670588f, 0.85098f, 0.913725f, 1.0f);
	cpuColorMap[4] = make_float4(0.878431f, 0.952941f, 0.972549f, 1.0f);
	cpuColorMap[5] = make_float4(1.0f, 1.0f, 0.74902f, 1.0f);
	cpuColorMap[6] = make_float4(0.996078f, 0.878431f, 0.564706f, 1.0f);
	cpuColorMap[7] = make_float4(0.992157f, 0.682353f, 0.380392f, 1.0f);
	cpuColorMap[8] = make_float4(0.956863f, 0.427451f, 0.262745f, 1.0f);
	cpuColorMap[9] = make_float4(0.843137f, 0.188235f, 0.152941f, 1.0f);
	cpuColorMap[10] = make_float4(0.647059f, 0.0f, 0.14902f, 1.0f);


	int width = BREWER_CTRL_PTS;
	int height = 0;


	// allocate array and copy colormap data
	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 32, 32, 32, cudaChannelFormatKindFloat);

	HANDLE_ERROR(cudaMallocArray(&gpuBrewer, &channelDesc, width, height));

	HANDLE_ERROR(cudaMemcpyToArray(gpuBrewer, 0, 0, cpuColorMap, sizeof(float4)*width, cudaMemcpyHostToDevice));

	// set texture parameters
    cudaTexBrewer.addressMode[0] = cudaAddressModeClamp;
	//texBrewer.addressMode[1] = cudaAddressModeClamp;
    cudaTexBrewer.filterMode = cudaFilterModeLinear;
    cudaTexBrewer.normalized = true;  // access with normalized texture coordinates

	// Bind the array to the texture
    HANDLE_ERROR(cudaBindTextureToArray( cudaTexBrewer, gpuBrewer, channelDesc));

}

static void destroyBrewer()
{
    HANDLE_ERROR(cudaFreeArray(gpuBrewer));

}

template<class T>
__global__ static void applyBrewer(T* gpuSource, unsigned char* gpuDest, unsigned int N, T minVal = 0, T maxVal = 1)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;
    if(i >= N) return;

	//compute the normalized value on [minVal maxVal]
	float a = (gpuSource[i] - minVal) / (maxVal - minVal);

	//lookup the color
	float shift = 1.0/BREWER_CTRL_PTS;
	float4 color = tex1D(cudaTexBrewer, a+shift);

	gpuDest[i * 3 + 0] = 255 * color.x;
	gpuDest[i * 3 + 1] = 255 * color.y;
	gpuDest[i * 3 + 2] = 255 * color.z;
}

template<class T>
__global__ static void applyGrayscale(T* gpuSource, unsigned char* gpuDest, unsigned int N, T minVal = 0, T maxVal = 1)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if(i >= N) return;

	//compute the normalized value on [minVal maxVal]
	float a = (gpuSource[i] - minVal) / (maxVal - minVal);

	gpuDest[i * 3 + 0] = 255 * a;
	gpuDest[i * 3 + 1] = 255 * a;
	gpuDest[i * 3 + 2] = 255 * a;
}

template<class T>
static void gpu2gpu(T* gpuSource, unsigned char* gpuDest, unsigned int nVals, T minVal = 0, T maxVal = 1, colormapType cm = cmGrayscale, int blockDim = 128)
{
	//This function converts a scalar field on the GPU to a color image on the GPU
	int gridDim = (nVals + blockDim - 1)/blockDim;
	if(cm == cmGrayscale)
		applyGrayscale<<<gridDim, blockDim>>>(gpuSource, gpuDest, nVals, minVal, maxVal);
	else if(cm == cmBrewer)
	{
		initBrewer();
		applyBrewer<<<gridDim, blockDim>>>(gpuSource, gpuDest, nVals, minVal, maxVal);
		destroyBrewer();
	}

}

template<class T>
static void gpu2cpu(T* gpuSource, unsigned char* cpuDest, unsigned int nVals, T minVal, T maxVal, colormapType cm = cmGrayscale)
{
    //this function converts a scalar field on the GPU to a color image on the CPU

    //first create the color image on the GPU

    //allocate GPU memory for the color image
    unsigned char* gpuDest;
    HANDLE_ERROR(cudaMalloc( (void**)&gpuDest, sizeof(unsigned char) * nVals * 3 ));

	//HANDLE_ERROR(cudaMemset(gpuSource, 0, sizeof(T) * nVals));

    //create the image on the gpu
    gpu2gpu(gpuSource, gpuDest, nVals, minVal, maxVal, cm);

	//HANDLE_ERROR(cudaMemset(gpuDest, 0, sizeof(unsigned char) * nVals * 3));

    //copy the image from the GPU to the CPU
    HANDLE_ERROR(cudaMemcpy(cpuDest, gpuDest, sizeof(unsigned char) * nVals * 3, cudaMemcpyDeviceToHost));

	HANDLE_ERROR(cudaFree( gpuDest ));

}

template<typename T>
static void gpu2image(T* gpuSource, std::string fileDest, unsigned int x_size, unsigned int y_size, T valMin, T valMax, colormapType cm = cmGrayscale)
{
	//allocate a color buffer
	unsigned char* cpuBuffer = (unsigned char*) malloc(sizeof(unsigned char) * 3 * x_size * y_size);

	//do the mapping
	gpu2cpu<T>(gpuSource, cpuBuffer, x_size * y_size, valMin, valMax, cm);

	//copy the buffer to an image
	buffer2image(cpuBuffer, fileDest, x_size, y_size);

	free(cpuBuffer);
}

#endif

template<class T>
static void cpu2cpu(T* cpuSource, unsigned char* cpuDest, unsigned int nVals, T valMin, T valMax, colormapType cm = cmGrayscale)
{
	int i;
	float a;
	float range = valMax - valMin;
	for(i = 0; i<nVals; i++)
	{
		//normalize to the range [valMin valMax]
		a = (cpuSource[i] - valMin) / range;

		cpuDest[i * 3 + 0] = 255 * a;
		cpuDest[i * 3 + 1] = 255 * a;
		cpuDest[i * 3 + 2] = 255 * a;
	}

}



template<typename T>
static void cpu2image(T* cpuSource, std::string fileDest, unsigned int x_size, unsigned int y_size, T valMin, T valMax, colormapType cm = cmGrayscale)
{
    //allocate a color buffer
	unsigned char* cpuBuffer = (unsigned char*) malloc(sizeof(unsigned char) * 3 * x_size * y_size);

	//do the mapping
	cpu2cpu<T>(cpuSource, cpuBuffer, x_size * y_size, valMin, valMax, cm);

	//copy the buffer to an image
	buffer2image(cpuBuffer, fileDest, x_size, y_size);

	free(cpuBuffer);

}

}}	//end namespace colormap and rts

#endif