/*======================================================================================

	Program:		RITK Plugins
	Module:			CUDASeparableNormalizedConvolutionImageFilterKernel.h
	Version:		$Revision: 898 $
	Date:			$Date: 2011-04-29 18:51:40 +0200 (Fri, 29 Apr 2011) $

	Copyright (c) 2010-2011
	Pattern Recognition Lab, 
    University Erlangen-Nuremberg, Germany
    All rights reserved.

	This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
	See RITKCopyright.txt for more information.

	http://www5.informatik.uni-erlangen.de/research/software/ritk

========================================================================================*/

#ifndef CUDASEPARABLENORMALIZEDCONVOLUTIONIMAGEFILTERKERNEL_H__
#define CUDASEPARABLENORMALIZEDCONVOLUTIONIMAGEFILTERKERNEL_H__

#include <cuda_runtime.h>
#include <driver_types.h>
#include <cuda_runtime_api.h>


// Division. If division remainder is neq zero then the result is ceiled
//----------------------------------------------------------------------------
#define DivUp(a,b) (((a) % (b) != 0) ? ((a)/(b) + 1) : ((a)/(b)))


// Kernel storage. Note the static length that defines our maximum kernel radius
//----------------------------------------------------------------------------
#define MAX_KERNEL_RADIUS 15
__constant__ float c_Kernel[2*MAX_KERNEL_RADIUS	+ 1];


// unsigned char to float conversion
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y>
__global__ void
UC2F_CastKernel(unsigned int NX, unsigned int NY, float *Output, size_t OutputPitch, const uchar4 *Input, size_t InputPitch)
{
	// Shared memory
	__shared__ uchar4 s_Data[BLOCKSIZE_Y][BLOCKSIZE_X];

	// Element offset
	const int InputBaseX = blockIdx.x * BLOCKSIZE_X + threadIdx.x;
    const int InputBaseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;
	const int OutputBaseX = blockIdx.x * BLOCKSIZE_X * 4 + threadIdx.x;
	const int OutputBaseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in y-direction
	if ( InputBaseY >= NY )
		return;

	// Set the data pointers
	Input += InputBaseY*InputPitch + InputBaseX;
	Output += OutputBaseY*OutputPitch + OutputBaseX;

	// Load the uc data to our shared mem
	if ( InputBaseX < NX/4 )
		s_Data[threadIdx.y][threadIdx.x] = Input[0];

	// Make sure all threads have pulled their data
	__syncthreads();

	// Store the result (4 floats)
	unsigned char *c = (unsigned char*)(&s_Data[threadIdx.y][0]);
	c += threadIdx.x;
	#pragma unroll
	for ( int i = 0; i < 4; i++ )
	{
		if ( OutputBaseX + i*BLOCKSIZE_X < NX )
			Output[i*BLOCKSIZE_X] = c[i*BLOCKSIZE_X];
	}
}


// float to unsigned char conversion
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y>
__global__ void
F2UC_CastKernel(unsigned int NX, unsigned int NY, uchar4 *Output, size_t OutputPitch, const float *Input, size_t InputPitch)
{
	// Shared memory
	__shared__ uchar4 s_Data[BLOCKSIZE_Y][BLOCKSIZE_X];

	// Element offset
	const int InputBaseX = blockIdx.x * BLOCKSIZE_X * 4 + threadIdx.x;
	const int InputBaseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;
	const int OutputBaseX = blockIdx.x * BLOCKSIZE_X + threadIdx.x;
	const int OutputBaseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in y-direction
	if ( InputBaseY >= NY )
		return;

	// Set the data pointers
	Input += InputBaseY*InputPitch + InputBaseX;
	Output += OutputBaseY*OutputPitch + OutputBaseX;

	// Load the data to our shared mem
	unsigned char *c = (unsigned char*)(&s_Data[threadIdx.y][0]);
	c += threadIdx.x;
	#pragma unroll
	for ( int i = 0; i < 4; i++ )
	{
		if ( InputBaseX + i*BLOCKSIZE_X < NX )
			c[i*BLOCKSIZE_X] = Input[i*BLOCKSIZE_X];
	}

	// Make sure all threads have pulled their data
	__syncthreads();

	// Store the result
	if ( OutputBaseX < NX/4 )
		Output[0] = s_Data[threadIdx.y][threadIdx.x];
}



// Element-wise multiplication
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y, int TILE_STEPS>
__global__ void
MultiplicationKernel(unsigned int NX, unsigned int NY, float *Output, const float *Input1, const float *Input2, size_t Pitch)
{
	// Element offset
	const int baseX = blockIdx.x * TILE_STEPS * BLOCKSIZE_X + threadIdx.x;
    const int baseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in y-direction
	if ( baseY >= NY )
		return;

	// Set our data pointers
	Input1 += baseY*Pitch + baseX;
	Input2 += baseY*Pitch + baseX;
	Output += baseY*Pitch + baseX;

	// Loop over the elements processed by this thread
	#pragma unroll
	for ( int i = 0; i < TILE_STEPS; i++ )
	{
		// Check for out-of-bounds in x-direction
		if ( baseX + i*BLOCKSIZE_X < NX )
		{
			Output[i*BLOCKSIZE_X] = Input1[i*BLOCKSIZE_X] * Input2[i*BLOCKSIZE_X];
		}
	}
}


// Row convolution
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y, int TILE_STEPS, int KERNEL_RADIUS>
__global__ void
CombinedRowConvolutionKernel(unsigned int NX, unsigned int NY, float *ImageDataOutput, float *MaskDataOutput, const float *ImageDataInput, const float *MaskDataInput, size_t Pitch)
{
	// First, determine the number of halos/aprons in terms of block sizes
	const int HALO_STEPS = DivUp(KERNEL_RADIUS, BLOCKSIZE_X);

	// Shared memory
	__shared__ float s_ImageData[BLOCKSIZE_Y][(TILE_STEPS + 2*HALO_STEPS) * BLOCKSIZE_X];
	__shared__ float s_MaskData[BLOCKSIZE_Y][(TILE_STEPS + 2*HALO_STEPS) * BLOCKSIZE_X];

	// Offset to the left halo edge
    const int baseX = (blockIdx.x * TILE_STEPS - HALO_STEPS) * BLOCKSIZE_X + threadIdx.x;
    const int baseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in y-direction
	if ( baseY >= NY )
		return;

	// Set our data pointers 
    ImageDataInput += baseY * Pitch + baseX;
	ImageDataOutput += baseY * Pitch + baseX;
	MaskDataInput += baseY * Pitch + baseX;
	MaskDataOutput += baseY * Pitch + baseX;

    // Load main data
    #pragma unroll
    for ( int i = HALO_STEPS; i < HALO_STEPS + TILE_STEPS; i++ )
	{
		if ( baseX + i*BLOCKSIZE_X < NX )
		{
			s_ImageData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = ImageDataInput[i*BLOCKSIZE_X];
			s_MaskData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = MaskDataInput[i*BLOCKSIZE_X];
		}
		else
		{
			s_ImageData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = 0;
			s_MaskData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = 0;
		}
	}

    // Load left halo
    #pragma unroll
    for ( int i = 0; i < HALO_STEPS; i++ )
	{
		s_ImageData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = ( baseX >= -i*BLOCKSIZE_X ) ? ImageDataInput[i*BLOCKSIZE_X] : 0;
		s_MaskData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = ( baseX >= -i*BLOCKSIZE_X ) ? MaskDataInput[i*BLOCKSIZE_X] : 0;
	}

    // Load right halo
    #pragma unroll
    for ( int i = HALO_STEPS + TILE_STEPS; i < HALO_STEPS + TILE_STEPS + HALO_STEPS; i++ )
	{
		s_ImageData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = ( NX-baseX > i*BLOCKSIZE_X ) ? ImageDataInput[i*BLOCKSIZE_X] : 0;
		s_MaskData[threadIdx.y][i*BLOCKSIZE_X + threadIdx.x] = ( NX-baseX > i*BLOCKSIZE_X ) ? MaskDataInput[i*BLOCKSIZE_X] : 0;
	}

	// Make sure all threads have pulled their data
	__syncthreads();

    // Loop over all pixels processed by this thread and compute/store their result
    #pragma unroll
    for ( int i = HALO_STEPS; i < HALO_STEPS + TILE_STEPS; i++ )
	{
		// Initialize our nominator and denominator. 
        float ImageDataNom = 0.f;
		float MaskDataNom = 0.f;
		float Denom = 0.f;

		// Convolution
        #pragma unroll
        for ( int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++ )
		{
			// Get the kernel and data values.
			float KernelValue = c_Kernel[KERNEL_RADIUS - j];
			float ImageDataValue = s_ImageData[threadIdx.y][threadIdx.x + i * BLOCKSIZE_X + j];
			float MaskDataValue = s_MaskData[threadIdx.y][threadIdx.x + i * BLOCKSIZE_X + j];

			// Update
			ImageDataNom += KernelValue*ImageDataValue;
			MaskDataNom += KernelValue*MaskDataValue;
			Denom += KernelValue;
		}

		// Store
		if ( baseX + i*BLOCKSIZE_X < NX )
		{
			ImageDataOutput[i*BLOCKSIZE_X] = ImageDataNom/Denom;
			MaskDataOutput[i*BLOCKSIZE_X] = MaskDataNom/Denom;
		}
    }
}


// Column convolution
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y,int TILE_STEPS, int KERNEL_RADIUS>
__global__ void
ColumnConvolutionKernel(unsigned int NX, unsigned int NY, float *Output, const float *Input, size_t Pitch)
{
	// First, determine the number of halos/aprons in terms of block sizes
	const int HALO_STEPS = DivUp(KERNEL_RADIUS, BLOCKSIZE_Y);

	// Shared memory
	__shared__ float s_Data[BLOCKSIZE_X][(TILE_STEPS + 2*HALO_STEPS) * BLOCKSIZE_Y + 1];

	// Offset to the upper left halo edge
    const int baseX = blockIdx.x * BLOCKSIZE_X + threadIdx.x;
    const int baseY = (blockIdx.y * TILE_STEPS - HALO_STEPS) * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in x-direction
	if ( baseX >= NX )
		return;

	// Set our data pointers 
    Input += baseY*Pitch + baseX;
	Output += baseY*Pitch + baseX;

	// Load main data
    #pragma unroll
    for ( int i = HALO_STEPS; i < HALO_STEPS + TILE_STEPS; i++ )
	{
		if ( baseY + i*BLOCKSIZE_Y < NY )
		{
			s_Data[threadIdx.x][i*BLOCKSIZE_Y + threadIdx.y] = Input[i*BLOCKSIZE_Y*Pitch];
		}
		else
		{
			s_Data[threadIdx.x][i*BLOCKSIZE_Y + threadIdx.y] = 0;
		}
	}

	// Load upper halo
    #pragma unroll
    for ( int i = 0; i < HALO_STEPS; i++ )
	{
		s_Data[threadIdx.x][i*BLOCKSIZE_Y + threadIdx.y] = ( baseY >= -i*BLOCKSIZE_Y ) ? Input[i*BLOCKSIZE_Y*Pitch] : 0;
	}


	// Load lower halo
    #pragma unroll
    for ( int i = HALO_STEPS + TILE_STEPS; i < HALO_STEPS + TILE_STEPS + HALO_STEPS; i++ )
	{
		s_Data[threadIdx.x][i*BLOCKSIZE_Y + threadIdx.y] = ( NY-baseY > i*BLOCKSIZE_Y ) ? Input[i*BLOCKSIZE_Y*Pitch] : 0;
	}

	// Make sure all threads have pulled their data
	__syncthreads();

	// Loop over all pixels processed by this thread and compute/store their result
    #pragma unroll
    for ( int i = HALO_STEPS; i < HALO_STEPS + TILE_STEPS; i++ )
	{
		// Initialize our nominator and denominator.
        float Nom = 0.f;
		float Denom = 0.f;

		// Convolution
        #pragma unroll
        for ( int j = -KERNEL_RADIUS; j <= KERNEL_RADIUS; j++ )
		{
			// Get the kernel and data values.
			float KernelValue = c_Kernel[KERNEL_RADIUS - j];
			float DataValue = s_Data[threadIdx.x][threadIdx.y + i * BLOCKSIZE_Y + j];

			// Update
			Nom += KernelValue*DataValue;
			Denom += KernelValue;
		}

		// Store
		if ( baseY + i*BLOCKSIZE_Y < NY )
			Output[i*BLOCKSIZE_Y*Pitch] = Nom/Denom;
    }
}


// Division
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y, int TILE_STEPS>
__global__ void
DivisionKernel(unsigned int NX, unsigned int NY, float *Output, const float *Nominator, const float *Denominator, size_t Pitch)
{
	// Element offset
	const int baseX = blockIdx.x * TILE_STEPS * BLOCKSIZE_X + threadIdx.x;
    const int baseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in y-direction
	if ( baseY >= NY )
		return;

	// Set our data pointers
	Nominator += baseY*Pitch + baseX;
	Denominator += baseY*Pitch + baseX;
	Output += baseY*Pitch + baseX;

	// Loop over the elements processed by this thread
	#pragma unroll
	for ( int i = 0; i < TILE_STEPS; i++ )
	{
		// Check for out-of-bounds in x-direction
		if ( baseX + i*BLOCKSIZE_X < NX )
		{
			// Pull values. Note that internally we're always operating with floats
			float Nom = Nominator[i*BLOCKSIZE_X];
			float Denom = Denominator[i*BLOCKSIZE_X];

			// Store the values
			if ( fabsf(Denom) > 1e-6f )
			{
				Output[i*BLOCKSIZE_X] = Nom/Denom;
			}
			else
			{
				Output[i*BLOCKSIZE_X] = 0.f;
			}
		}
	}
}


// Final assembly kernel
//----------------------------------------------------------------------------
template<int BLOCKSIZE_X, int BLOCKSIZE_Y, int TILE_STEPS>
__global__ void
FinalAssemblyKernel
(
	unsigned int NX,
	unsigned int NY,
	float *ImageDataOutput,
	float *MaskDataOutput,
	const float *ImageDataInput,
	const float *ImageDataFiltered,
	const float *MaskDataInput,
	size_t Pitch
)
{
	// Element offset
	const int baseX = blockIdx.x * BLOCKSIZE_X * TILE_STEPS + threadIdx.x;
    const int baseY = blockIdx.y * BLOCKSIZE_Y + threadIdx.y;

	// Early termination due to out-of-bounds in y-direction
	if ( baseY >= NY )
		return;

	// Set our data pointers
	ImageDataOutput += baseY*Pitch + baseX;
	MaskDataOutput += baseY*Pitch + baseX;
	ImageDataInput += baseY*Pitch + baseX;
	ImageDataFiltered += baseY*Pitch + baseX;
	MaskDataInput += baseY*Pitch + baseX;
	
	// Loop over the elements processed by this thread
	#pragma unroll
	for ( int i = 0; i < TILE_STEPS; i++ )
	{
		// Check for out-of-bounds in x-direction
		if ( baseX + i*BLOCKSIZE_X < NX )
		{
			// Pull values.
			float _ImageDataInput = ImageDataInput[i*BLOCKSIZE_X];
			float _ImageDataFiltered = ImageDataFiltered[i*BLOCKSIZE_X];
			unsigned char _MaskDataInput = MaskDataInput[i*BLOCKSIZE_X];

			// Keep the original value if we had a valid flag
			if ( _MaskDataInput == 1 )
			{
				ImageDataOutput[i*BLOCKSIZE_X] = _ImageDataInput;
				MaskDataOutput[i*BLOCKSIZE_X] = _MaskDataInput;
			}
			else
			{
				ImageDataOutput[i*BLOCKSIZE_X] = _ImageDataFiltered;
				MaskDataOutput[i*BLOCKSIZE_X] = _ImageDataFiltered == 0 ? 0.f : 1.f;
			}
		}
	}
}


#endif // CUDASEPARABLENORMALIZEDCONVOLUTIONIMAGEFILTERKERNEL_H__