#include "ritkCudaContext.h"
#include "ritkCudaRegularMemoryImportImageContainer.h"

#include "KinectImageCudaDepthToWorldFilter.h"

// CUDA interface
//----------------------------------------------------------------------------
extern "C"
void
CUDARangeToWorldKernel(const cudaArray *InputImageArray, float3 *DeviceOutput, int w, int h, float X2Z, float Y2Z);


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::KinectImageCudaDepthToWorldFilter()
{
	m_XToZ = 0;
	m_YToZ = 0;

	// Initialize CUDA memory members
	m_InputArray = 0;
	m_CurrentSize[0] = 0;
	m_CurrentSize[1] = 0;
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::~KinectImageCudaDepthToWorldFilter()
{
	CudaBlockedMethodCallMacro(FreeCudaMemory);
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::SetConversionParameters(float XToZ, float YToZ)
{
	m_XToZ = XToZ;
	m_YToZ = YToZ;
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::GetConversionParameters(float *XToZ, float *YToZ)
{
	*XToZ = m_XToZ;
	*YToZ = m_YToZ;
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::AllocateCudaMemory()
{
	// Size in bytes of the image
	typename InputImageType::SizeType Size = this->GetInput()->GetBufferedRegion().GetSize();

	// Allocate CUDA device memory due to changed regions
	if ( m_CurrentSize != Size )
	{
		// First clean up old memory
		FreeCudaMemory();

		// Allocate new memory and set the member that hold the supported size
		cudaChannelFormatDesc ImageChannelDesc = cudaCreateChannelDesc(32,0,0,0,cudaChannelFormatKindFloat);
		ritkCudaSafeCall(cudaMallocArray(&m_InputArray,&ImageChannelDesc,this->GetInput()->GetBufferedRegion().GetSize()[0],this->GetInput()->GetBufferedRegion().GetSize()[1]));

		// A CUDA output container. Workaround due to missing virtual modifiers in itk
		ritk::CudaRegularMemoryImportImageContainerF3::Pointer Container = ritk::CudaRegularMemoryImportImageContainerF3::New();
		ritk::CudaRegularMemoryImportImageContainerF3::SizeType ContainerSize;
		ContainerSize[0] = this->GetOutput()->GetPixelContainer()->Capacity();
		Container->SetContainerSize(ContainerSize);
		Container->Reserve(this->GetOutput()->GetPixelContainer()->Capacity());
		this->GetOutput()->SetPixelContainer(Container);

		// Update the size member
		m_CurrentSize = Size;
	}	

	// Set new cuda memory container (possibly due to a DisconnectPipeline)
	if ( !dynamic_cast<ritk::CudaRegularMemoryImportImageContainerF3*>(this->GetOutput()->GetPixelContainer()) )
	{
		ritk::CudaRegularMemoryImportImageContainerF3::Pointer Container = ritk::CudaRegularMemoryImportImageContainerF3::New();
		Container->Reserve(this->GetOutput()->GetPixelContainer()->Capacity());
		this->GetOutput()->SetPixelContainer(Container);
	}
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::FreeCudaMemory()
{
	// Free device memory
	if ( m_InputArray )
		ritkCudaSafeCall(cudaFreeArray(m_InputArray));

	// Update the size member
	m_CurrentSize[0] = 0;
	m_CurrentSize[1] = 0;
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::DoFiltering()
{
	// Copy data
	ritk::CudaRegularMemoryImportImageContainerF::ConstPointer CudaContainer = dynamic_cast<const ritk::CudaRegularMemoryImportImageContainerF*>(this->GetInput()->GetPixelContainer());
	if ( CudaContainer )
	{
		//std::cout << "Found a cuda container..." << std::endl;
		ritkCudaSafeCall(cudaMemcpyToArray(m_InputArray, 0, 0, CudaContainer->GetCudaMemoryPointer(), m_CurrentSize[0]*m_CurrentSize[1] * sizeof(float), cudaMemcpyDeviceToDevice));
	}
	else
	{
		// Copy data to device
		ritkCudaSafeCall(cudaMemcpyToArray(m_InputArray, 0, 0, this->GetInput()->GetBufferPointer(), m_CurrentSize[0]*m_CurrentSize[1] * sizeof(float), cudaMemcpyHostToDevice));
	}


	// Copy the input data to the device
	//cutilSafeCall(cudaMemcpyToArray(m_InputArray, 0, 0, GetInput()->GetBufferPointer(), m_CurrentSize[0]*m_CurrentSize[1] * sizeof(float), cudaMemcpyHostToDevice));

	// Get the output memory
	ritk::CudaRegularMemoryImportImageContainerF3::Pointer OutputCudaContainer = dynamic_cast<ritk::CudaRegularMemoryImportImageContainerF3*>(this->GetOutput()->GetPixelContainer());
	if ( !OutputCudaContainer )
		itkExceptionMacro(<<"Invalid pixel container detected!");
	float3 *Output = (float3*)OutputCudaContainer->GetCudaMemoryPointer();

	// Kernel invocation
	CUDARangeToWorldKernel(m_InputArray, Output, m_CurrentSize[0], m_CurrentSize[1], m_XToZ, m_YToZ);

	// Copy data to the host
	//ritkCudaSafeCall(cudaMemcpy(GetOutput()->GetBufferPointer(), Output, sizeof(float3)*m_CurrentSize[0]*m_CurrentSize[1], cudaMemcpyDeviceToHost));
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
KinectImageCudaDepthToWorldFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::GenerateData()
{
	// Sanity checks
	if ( !this->GetInput() )
		itkExceptionMacro(<<"No input provided!");
	if ( fabs(m_XToZ) < 1e-3 || fabs(m_YToZ) < 1e-3 )
		itkExceptionMacro(<<"Invalid conversion parameters detected!");

	// Ensure that we have a valid output
	this->GetOutput()->SetRequestedRegion(this->GetInput()->GetRequestedRegion());
	this->GetOutput()->SetBufferedRegion(this->GetInput()->GetBufferedRegion());
	this->GetOutput()->SetLargestPossibleRegion(this->GetInput()->GetLargestPossibleRegion());
	this->GetOutput()->Allocate();

	// Allocate device memory
	CudaBlockedMethodCallMacro(AllocateCudaMemory);

	// Do the actual filtering
	CudaBlockedMethodCallMacro(DoFiltering);
}

