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

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkTruncatedSignedDistanceFunctionRayCaster.txx
	Version:		$Revision: 3136 $
	Date:			$Date: 2012-04-19 14:42:59 +0200 (Do, 19 Apr 2012) $

	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

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

#include "ritkTruncatedSignedDistanceFunctionRayCaster.h"
#include "ritkTruncatedSignedDistanceFunctionRayCasterInterfaces.h"
#include "ritkCudaRayCasterInterfaces.h"
#include "ritkCudaPitchedMemoryImportImageContainer.h"

#include <vnl/vnl_double_3.h>

#include <itkRescaleIntensityImageFilter.h>

//#define PERFORMANCE_CHECK

namespace ritk
{

//----------------------------------------------------------------------------
template<typename TRImage>
TruncatedSignedDistanceFunctionRayCaster<TRImage>::TruncatedSignedDistanceFunctionRayCaster()
{
	// Default size
	Superclass::m_OutputSize[0] = 512;
	Superclass::m_OutputSize[1] = 512;
}


//----------------------------------------------------------------------------
template<typename TRImage>
TruncatedSignedDistanceFunctionRayCaster<TRImage>::~TruncatedSignedDistanceFunctionRayCaster()
{
}


//----------------------------------------------------------------------------
template<typename TRImage>
typename TruncatedSignedDistanceFunctionRayCaster<TRImage>::DataObjectPointer
TruncatedSignedDistanceFunctionRayCaster<TRImage>::MakeOutput(unsigned int idx)
{
	return Superclass::MakeOutput(idx);
}


//----------------------------------------------------------------------------
template<typename TRImage>
void
TruncatedSignedDistanceFunctionRayCaster<TRImage>::GenerateOutputInformation()
{
	// Different dimensions are handles by the ray caster base class
	Superclass::GenerateOutputInformation();

	// Set the correct range output container.
	typedef ritk::CudaPitchedMemoryImportImageContainer<unsigned long, typename OutputRangeImageType::RangeImageType::PixelType> OutputRangeDataCudaContainerType;
	typename OutputRangeDataCudaContainerType::Pointer OutputRangeDataCudaContainer = dynamic_cast<OutputRangeDataCudaContainerType*>(this->GetOutput()->GetRangeImage()->GetPixelContainer());
	if ( !OutputRangeDataCudaContainer )
	{
		OutputRangeDataCudaContainer = OutputRangeDataCudaContainerType::New();
		this->GetOutput()->GetRangeImage()->SetPixelContainer(OutputRangeDataCudaContainer);
	}
	OutputRangeDataCudaContainer->SetContainerSize(Superclass::m_OutputSize);

	// Set the correct world coord output container
	/*typedef ritk::CudaPitchedMemoryImportImageContainer<unsigned long, typename OutputRangeImageType::WorldCoordImageType::PixelType> OutputWorldCoordDataCudaContainerType;
	OutputWorldCoordDataCudaContainerType::Pointer OutputWorldCoordDataCudaContainer = dynamic_cast<OutputWorldCoordDataCudaContainerType*>(this->GetOutput()->GetWorldCoordImage()->GetPixelContainer());
	if ( !OutputWorldCoordDataCudaContainer )
	{
		OutputWorldCoordDataCudaContainer = OutputWorldCoordDataCudaContainerType::New();
		this->GetOutput()->GetWorldCoordImage()->SetPixelContainer(OutputWorldCoordDataCudaContainer);
	}
	OutputWorldCoordDataCudaContainer->SetContainerSize(Superclass::m_OutputSize);*/
}


//----------------------------------------------------------------------------
template<typename TRImage>
void
TruncatedSignedDistanceFunctionRayCaster<TRImage>::GenerateData()
{
#ifdef PERFORMANCE_CHECK
	cudaEvent_t start_event, stop_event;
    cudaEventCreateWithFlags(&start_event, cudaEventBlockingSync);
    cudaEventCreateWithFlags(&stop_event, cudaEventBlockingSync);
	cudaEventRecord(start_event, 0);
	int NumRuns = 50;
	for ( int i = 0; i < NumRuns; i++ )
	{
#endif

	// Allocate outputs as a first step
	this->AllocateOutputs();
	this->GetOutput()->GetWorldCoordImage()->Allocate();

	// Set the parameters (adjust to your needs)
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::FX] = 5.9421434211923247e+02;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::FY] = 5.9104053696870778e+02;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::CX] = 3.3930780975300314e+02;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::CY] = 2.4273913761751615e+02;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::K1] = -2.6386489753128833e-01;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::K2] = 9.9966832163729757e-01;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::K3] = -1.3053628089976321e+00;

	// Sanity check for correct input cuda container
	typedef ritk::Cuda3DArrayImportImageContainer<unsigned long, typename InputTSDFType::PixelType> InputCudaContainerType;
	typename InputCudaContainerType::ConstPointer InputCudaContainer = dynamic_cast<const InputCudaContainerType*>(this->GetInput()->GetPixelContainer());
	if ( !InputCudaContainer )
	{
		itkExceptionMacro(<<"Invalid input image container detected. Must be an ritk::Cuda3DArrayImportImageContainer<unsigned long, float>");
	}

	// Sanity check for correct range data output container
	typedef ritk::CudaPitchedMemoryImportImageContainer<unsigned long, typename OutputRangeImageType::RangeImageType::PixelType> OutputRangeDataCudaContainerType;
	typename OutputRangeDataCudaContainerType::Pointer OutputRangeDataCudaContainer = dynamic_cast<OutputRangeDataCudaContainerType*>(this->GetOutput()->GetRangeImage()->GetPixelContainer());
	if ( !OutputRangeDataCudaContainer )
	{
		itkExceptionMacro(<<"Invalid range data container detected. Must be an ritk::CudaPitchedMemoryImportImageContainer<unsigned long, float>. Did you disconnect this filter's output?");
	}

	// Set the correct world coord output container
	/*typedef ritk::CudaPitchedMemoryImportImageContainer<unsigned long, typename OutputRangeImageType::WorldCoordImageType::PixelType> OutputWorldCoordDataCudaContainerType;
	OutputWorldCoordDataCudaContainerType::Pointer OutputWorldCoordDataCudaContainer = dynamic_cast<OutputWorldCoordDataCudaContainerType*>(this->GetOutput()->GetWorldCoordImage()->GetPixelContainer());
	if ( !OutputWorldCoordDataCudaContainer )
	{
		itkExceptionMacro(<<"Invalid world coord data container detected. Must be an ritk::CudaPitchedMemoryImportImageContainer<unsigned long, itk::Vector<float,3> >. Did you disconnect this filter's output?");
	}*/


	// Set up the volume parameters
	InputTSDFConstPointer InputImage = this->GetInput();
	ritkCudaRayCasterVolumeParameters VP;
	VP.Size.x = InputImage->GetRequestedRegion().GetSize()[0];
	VP.Size.y = InputImage->GetRequestedRegion().GetSize()[1];
	VP.Size.z = InputImage->GetRequestedRegion().GetSize()[2];
	VP.Origin.x = InputImage->GetOrigin()[0];
	VP.Origin.y = InputImage->GetOrigin()[1];
	VP.Origin.z = InputImage->GetOrigin()[2];
	VP.Spacing.x = InputImage->GetSpacing()[0];
	VP.Spacing.y = InputImage->GetSpacing()[1];
	VP.Spacing.z = InputImage->GetSpacing()[2];
	VP.Center.x = 0.5*InputImage->GetRequestedRegion().GetSize()[0]*InputImage->GetSpacing()[0];
	VP.Center.y = 0.5*InputImage->GetRequestedRegion().GetSize()[1]*InputImage->GetSpacing()[1];
	VP.Center.z = 0.5*InputImage->GetRequestedRegion().GetSize()[2]*InputImage->GetSpacing()[2];
	VP.ChannelDesc = &InputCudaContainer->GetChannelDesc();
	VP.DataArray = InputCudaContainer->GetCudaMemoryArray();

	// Set up the range image parameters. Note that the pitch is in elements NOT bytes
	typename OutputRangeImageType::RangeImagePointer OutputRangeImage = this->GetOutput()->GetRangeImage();
	ritkCudaRayCasterProjectionImageParameters<typename OutputRangeImageType::RangeImageType::PixelType> RIP;
	RIP.Size.x = OutputRangeImage->GetRequestedRegion().GetSize()[0];
	RIP.Size.y = OutputRangeImage->GetRequestedRegion().GetSize()[1];
	RIP.Origin.x = OutputRangeImage->GetOrigin()[0];
	RIP.Origin.y = OutputRangeImage->GetOrigin()[1];
	RIP.Spacing.x = OutputRangeImage->GetSpacing()[0];
	RIP.Spacing.y = OutputRangeImage->GetSpacing()[1];
	RIP.Pitch = OutputRangeDataCudaContainer->GetPitch()/sizeof(typename OutputRangeDataCudaContainerType::Element);
	RIP.DataPtr = OutputRangeDataCudaContainer->GetCudaMemoryPointer();

	// Set up the normal image parameters. Note that the pitch is in elements NOT bytes
	typename OutputRangeImageType::WorldCoordImagePointer OutputWorldCoordImage = this->GetOutput()->GetWorldCoordImage();
	ritkCudaRayCasterProjectionImageParameters<float3> WIP;
	WIP.Size.x = OutputWorldCoordImage->GetRequestedRegion().GetSize()[0];
	WIP.Size.y = OutputWorldCoordImage->GetRequestedRegion().GetSize()[1];
	WIP.Origin.x = OutputWorldCoordImage->GetOrigin()[0];
	WIP.Origin.y = OutputWorldCoordImage->GetOrigin()[1];
	WIP.Spacing.x = OutputWorldCoordImage->GetSpacing()[0];
	WIP.Spacing.y = OutputWorldCoordImage->GetSpacing()[1];
	/*WIP.Pitch = OutputWorldCoordDataCudaContainer->GetPitch()/sizeof(OutputWorldCoordDataCudaContainerType::Element);
	WIP.DataPtr = (float3*)OutputWorldCoordDataCudaContainer->GetCudaMemoryPointer();*/
	WIP.Pitch = 0;
	WIP.DataPtr = NULL;

	// Set up the projection parameters
	ritkCudaRayCasterProjectionParameters PP;
	// Inverse projection matrix
	ProjectionMatrixType P = this->GetProjectionMatrixInput()->Get();
	vnl_double_3x3 tmp;
	tmp(0,0) = P(0,0); tmp(0,1) = P(0,1); tmp(0,2) = P(0,2);
	tmp(1,0) = P(1,0); tmp(1,1) = P(1,1); tmp(1,2) = P(1,2);
	tmp(2,0) = P(2,0); tmp(2,1) = P(2,1); tmp(2,2) = P(2,2);
	vnl_svd<double> SVD(tmp);
	ProjectionMatrixInverseType IPM = SVD.pinverse();
	for ( unsigned int row = 0; row < 3; row++ )
		for ( unsigned int col = 0; col < 3; col++ )
			PP.InverseProjectionMatrix[row*3+col] = IPM(row,col);
	// Extract the translation part out of the projection matrix
	vnl_double_3 q;
	q(0) = -1*P(0,3);
	q(1) = -1*P(1,3);
	q(2) = -1*P(2,3);
	q = IPM*q;
	PP.OpticalCenter.x = q[0];
	PP.OpticalCenter.y = q[1];
	PP.OpticalCenter.z = q[2];


	// Compute ray direction
	/*unsigned int CenterX = 320;
	unsigned int CenterY = 240;
	float RayX = PP.InverseProjectionMatrix[0]*CenterX + PP.InverseProjectionMatrix[1]*CenterY + PP.InverseProjectionMatrix[2];
	float RayY = PP.InverseProjectionMatrix[3]*CenterX + PP.InverseProjectionMatrix[4]*CenterY + PP.InverseProjectionMatrix[5];
	float RayZ = PP.InverseProjectionMatrix[6]*CenterX + PP.InverseProjectionMatrix[7]*CenterY + PP.InverseProjectionMatrix[8];
	float Mag = -1.f/sqrt((RayX)*(RayX) + (RayY)*(RayY) + (RayZ)*(RayZ));
	RayX *= Mag;
	RayY *= Mag;
	RayZ *= Mag;
	std::cout << "Center ray direction = [" << RayX << "," << RayY << "," << RayZ << "]" << std::endl;*/

	// Call the kernel
	ritkTSDFRayCaster(VP,RIP,WIP,PP, this->GetInput()->GetCutoffValue(), this->GetInput()->GetInvalidValue());
	
	OutputRangeDataCudaContainer->SynchronizeHost();
	this->GetOutput()->Finish();

	// Synchronize the host
	this->GetOutput()->SynchronizeGenericImportImageContainers();

	/*std::ofstream out("D:/Temp/WorldCoordTest.txt");
	for ( unsigned int i = 0; i < WIP.Size.x*WIP.Size.y; i++ )
	{
		out << this->GetOutput()->GetWorldCoordImage()->GetBufferPointer()[i] << std::endl;
	}*/


	/*typedef itk::Image<float,2> OutputImageType;
	typedef itk::RescaleIntensityImageFilter<OutputRangeImageType::RangeImageType, OutputImageType> RescalerType;
	RescalerType::Pointer Rescaler = RescalerType::New();
	Rescaler->SetOutputMinimum(0);
	Rescaler->SetOutputMaximum(1);
	Rescaler->SetInput(this->GetOutput()->GetRangeImage());
	Rescaler->Update();
	std::cout << "Min=" << Rescaler->GetInputMinimum() << " Max=" << Rescaler->GetInputMaximum() << std::endl;*/


	// Update the flag image
	float *RangePtr = this->GetOutput()->GetRangeImage()->GetBufferPointer();
	unsigned char *FlagPtr = this->GetOutput()->GetFlagImage()->GetBufferPointer();
	unsigned long NumPixels = this->GetOutput()->GetRequestedRegion().GetNumberOfPixels();
	for ( unsigned long l = 0; l < NumPixels; l++ )
	{
		float f = *RangePtr;
		if ( abs(f) > 1e-3 )
			*FlagPtr = OutputRangeImageType::VALID;
		else
			*FlagPtr = !OutputRangeImageType::VALID;
		
		RangePtr++;
		FlagPtr++;
	}

#ifdef PERFORMANCE_CHECK
	}
	cudaEventRecord(stop_event, 0);
    cudaEventSynchronize(stop_event);   
	float Time = 0;
    cudaEventElapsedTime(&Time, start_event, stop_event);
	std::cout << "Elapsed time for predicting surface: " << (Time/NumRuns) << std::endl;
#endif
}




} // ritk
