/// Copyright (c) 2010 Ruben Penalva. All rights reserved.
/// Redistribution and use in source and binary forms, with or without modification, 
/// are permitted provided that the following conditions are met:
///    1. Redistributions of source code must retain the above copyright notice, this 
///    list of conditions and the following disclaimer.
///    2. Redistributions in binary form must reproduce the above copyright notice,
///    this list of conditions and the following disclaimer in the documentation and/or
///    other materials provided with the distribution.
///    3. The name of the author may not be used to endorse or promote products derived
///    from this software without specific prior written permission.
///
/// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,
/// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 
/// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, 
/// INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED 
/// TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
/// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 
/// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 
/// THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

/// @file kernel.cu
/// This source file is part of CUDA Ray Casting test.
/// The file holds the CUDA Raycasting kernel.
/// @author Ruben Penalva (ruben.penalva@rpenalva.com)

#ifndef _KERNEL_CU_
#define _KERNEL_CU_

#include <cuda_runtime.h>
#include <device_functions.h>
#include <math.h>

texture<unsigned char,  3, cudaReadModeNormalizedFloat> volumeTexture;
texture<float4, 1, cudaReadModeElementType> transferTexture;

inline __device__ float3 operator*(float3 v, float s)
{
	return make_float3(v.x * s, v.y * s, v.z * s);
}

inline __device__ float3 operator+(float3 v, float s)
{
	return make_float3(v.x + s, v.y + s, v.z + s);
}

inline __device__ float Length(float3 v)
{
	return sqrtf(v.x * v.x + v.y * v.y + v.z * v.z);
}

inline __device__ float InvLength(float3 v)
{
	return 1.0f / Length(v);
}

inline __device__ float3 Normalize(float3 v)
{
	return v * InvLength(v);
}

inline __device__ float3 operator*(float* matrix, float3 v)
{
	float x = matrix[0] * v.x + matrix[1] * v.y + matrix[2] * v.z;
	float y = matrix[4] * v.x + matrix[5] * v.y + matrix[6] * v.z;
	float z = matrix[8] * v.x + matrix[9] * v.y + matrix[10] * v.z;

	return make_float3(x, y, z);
}

inline __device__ float3 operator-(float3 lv, float3 rv)
{
	return make_float3(lv.x - rv.x, lv.y - rv.y, lv.z - rv.z);
}

inline __device__ float3 operator*(float3 lv, float3 rv)
{
	return make_float3(lv.x * rv.x, lv.y * rv.y, lv.z * rv.z);
}

inline __device__ float3 operator/(float3 lv, float3 rv)
{
	return make_float3(lv.x / rv.x, lv.y / rv.y, lv.z / rv.z);
}

inline __device__ float3 operator+(float3 lv, float3 rv)
{
	return make_float3(lv.x + rv.x, lv.y + rv.y, lv.z + rv.z);
}


inline __device__ float4 operator+(float4 lv, float4 rv)
{
	return make_float4(lv.x + rv.x, lv.y + rv.y, lv.z + rv.z, lv.w + rv.w);
}

__inline__ __device__ void Swap(float* t1, float* t2)
{
	float t = *t2;
	*t2 = *t1;
	*t1 = t;
}

__inline__ __device__ void SwapGR(float3* t1, float3* t2)
{
	if (t1->x > t2->x)
	{
		Swap(&t1->x, &t2->x);
	}

	if (t1->y > t2->y)
	{
		Swap(&t1->y, &t2->y);
	}

	if (t1->z > t2->z)
	{
		Swap(&t1->z, &t2->z);
	}
}

__inline__ __device__ float Max(float3 t1)
{
	return fmaxf(fmaxf(t1.x, t1.y), t1.z);
}

__inline__ __device__ float Min(float3 t1)
{
	return fminf(fminf(t1.x, t1.y), t1.z);
}

/// This is a really straightforward and really slow implementation of Ericcson book intersectionrayaabb function
/// To get a better idea of a quick implementation of ray-aabox intersection check volumerender example within
/// CUDA SDK Toolkit
inline __device__ int IntersectRayAABB(float3 p, float3 d, float3 abboxMin, float3 abboxMax, float* tmin, float* tmax)
{
	float3 ood = make_float3(1.0f, 1.0f, 1.0f) / d;
	float3 t1 = (abboxMin - p) * ood;
	float3 t2 = (abboxMax - p) * ood;

	SwapGR(&t1, &t2);

	*tmin = Max(t1);
	*tmax = Min(t2);

	return *tmax > *tmin;
}

__global__ void RayCast(unsigned int* pixelBuffer, int pbWidth, int pbHeight, 
									float* cameraPos, float* invViewMatrix, float* invVolMatrix,
									float* bboxMin, float* bboxMax,
									int steps, float step)
{
	// 2D Thread ID
	unsigned int x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
	unsigned int y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;
	float u = (x / (float) pbWidth) * 2.0f - 1.0f;
	float v = (y / (float) pbHeight) * 2.0f - 1.0f;

	float3 origin = make_float3(cameraPos[0], cameraPos[1], cameraPos[2]);
	float3 direction = invViewMatrix * Normalize(make_float3(u, v, -3.0f));
	direction = invVolMatrix * direction;

	float tnear;
	float tfar;
	int intersected = IntersectRayAABB(origin, direction, 
									   make_float3(bboxMin[0], bboxMin[1], bboxMin[2]), 
									   make_float3(bboxMax[0], bboxMax[1], bboxMax[2]),
									   &tnear, &tfar);
	
	const float3 bgColor = {0.5f, 0.5f, 0.7f};
	float4 finalColor = {0.0f, 0.0f, 0.0f, 0.0f};
	if (!intersected)
	{
		finalColor.x = bgColor.x;
		finalColor.y = bgColor.y;
		finalColor.z = bgColor.z;	
	}
	else
	{
		if (tnear < 0.0f)
		{
			tnear = 0.0f;
		}

		float3 position = origin;

		float t = tnear;
		for (int i = 0; i < steps; ++i)
		{
			float3 position = origin + direction * t;
			position = position + 0.5f;

			float density = tex3D(volumeTexture, position.x, position.y, position.z);

			float4 color = tex1D(transferTexture, density);
			color.x *= color.w;
			color.y *= color.w;
			color.z *= color.w;

			finalColor = finalColor + color;

			if (finalColor.w > 1.0f || t > tfar)
			{
				break;
			}

			t += step;
		}

		finalColor.x = finalColor.x * finalColor.w + bgColor.x * (1 - finalColor.w);
		finalColor.y = finalColor.y * finalColor.w + bgColor.y * (1 - finalColor.w);
		finalColor.z = finalColor.z * finalColor.w + bgColor.z * (1 - finalColor.w);
	}

	unsigned int i = __umul24(y, pbWidth) + x;
	pixelBuffer[i] = (unsigned int(1.0f * 255) << 24) | 
					 (unsigned int(__saturatef(finalColor.z) * 255) << 16) |  
					 (unsigned int(__saturatef(finalColor.y) * 255) << 8)  | 
					  unsigned int(__saturatef(finalColor.x) * 255);
}

extern "C"
void CallCUDAKernel(dim3 gridDim, dim3 blockDim, 
				unsigned int* pixelBuffer, int pbWidth, int pbHeight, 
				float* cameraPos, float* invViewMatrix, float* invVolMatrix, 
				float* bboxMin, float* bboxMax, int steps, float step)
{
	RayCast<<<gridDim, blockDim>>>(pixelBuffer, pbWidth, pbHeight, 
											   cameraPos, invViewMatrix, invVolMatrix, 
											   bboxMin, bboxMax, steps, step);
}

cudaArray* d_volumeArray;
cudaArray* d_transferFuncArray;

extern "C"
void CreateCUDAVolume(int texWidth, int texHeight, int texDepth, unsigned char* dataPtr)
{
	cudaExtent volumeSize = {texWidth, texHeight, texDepth};
	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<unsigned char>();
	cudaMalloc3DArray(&d_volumeArray, &channelDesc, volumeSize);

	cudaMemcpy3DParms copyParams = {0};
	copyParams.srcPtr   = make_cudaPitchedPtr(dataPtr, texWidth * sizeof(unsigned char), texWidth, texHeight);
	copyParams.dstArray = d_volumeArray;
	copyParams.extent   = volumeSize;
	copyParams.kind     = cudaMemcpyHostToDevice;
	cudaMemcpy3D(&copyParams);  

	volumeTexture.normalized = true;
	volumeTexture.filterMode = cudaFilterModeLinear;
	volumeTexture.addressMode[0] = cudaAddressModeClamp;
	volumeTexture.addressMode[1] = cudaAddressModeClamp;

	cudaBindTextureToArray(volumeTexture, d_volumeArray, channelDesc);
}

extern "C"
void CreateCUDATransferFunction(float* transferFunctionPtr, int transferFunctionWidth)
{
	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float4>();
	cudaMallocArray(&d_transferFuncArray, &channelDesc, transferFunctionWidth, 1); 
	cudaMemcpyToArray(d_transferFuncArray, 0, 0, transferFunctionPtr, transferFunctionWidth * sizeof(float4), cudaMemcpyHostToDevice);

	transferTexture.filterMode = cudaFilterModeLinear;
	transferTexture.normalized = true;
	transferTexture.addressMode[0] = cudaAddressModeClamp;

	cudaBindTextureToArray(transferTexture, d_transferFuncArray, channelDesc);
}

extern "C"
void UpdateCUDATransferFunction(float* transferFunctionPtr, int transferFunctionWidth)
{
	cudaMemcpyToArray(d_transferFuncArray, 0, 0, transferFunctionPtr, transferFunctionWidth * sizeof(float4), cudaMemcpyHostToDevice);
}

extern "C"
void DestroyCudaVolume(void)
{
	cudaFreeArray(d_volumeArray);
}

extern "C"
void DestroyCUDATransferFunction(void)
{
	cudaFreeArray(d_transferFuncArray);
}
#endif