#include "CUDABilateralImageFilter.h"

#include "ritkCudaRegularMemoryImportImageContainer.h"


// Cuda interface
//----------------------------------------------------------------------------
extern "C"
float
BilateralFilter(const cudaArray *InputImageArray, const cudaArray *PrecomputedGaussianArray, float *DeviceOutput, size_t OutputPitch, int w, int h, int KernelRadius, float SigmaD, float SigmaR);


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
CUDABilateralImageFilter<TInput,TOutput>
::CUDABilateralImageFilter()
{
	// Default kernel and sigmas
	m_KernelSize = 5;
	m_SpatialSigma = 50;
	m_RangeSigma = 50;

	// Invalid CUDA memory by default
	m_CurrentSize[0] = 0;
	m_CurrentSize[1] = 0;
	m_OutputImageContainer = NULL;
	m_InputImageContainer = NULL;

	// PreComputed gaussian
	m_PreComputedGaussianContainer = NULL;
	m_PreComputedGaussianKernelSize = 0;
	m_PreComputedGaussianSigma = 0;
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
CUDABilateralImageFilter<TInput,TOutput>
::~CUDABilateralImageFilter()
{
	// CUDA memory should be cleaned up automatically
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
void
CUDABilateralImageFilter<TInput,TOutput>
::SetSpatialSigma(float SpatialSigma)
{
	if ( SpatialSigma == m_SpatialSigma )
		return;

	m_SpatialSigma = SpatialSigma;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
void
CUDABilateralImageFilter<TInput,TOutput>
::SetRangeSigma(float RangeSigma)
{
	if ( RangeSigma == m_RangeSigma )
		return;

	m_RangeSigma = RangeSigma;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
void
CUDABilateralImageFilter<TInput,TOutput>
::SetKernelSize(int size)
{
	m_KernelSize = size;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
float
CUDABilateralImageFilter<TInput,TOutput>
::GetSpatialSigma()
{
	return m_SpatialSigma;
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
float
CUDABilateralImageFilter<TInput,TOutput>
::GetRangeSigma()
{
	return m_RangeSigma;
}


//----------------------------------------------------------------------------
template<class TInput, class TOutput>
void
CUDABilateralImageFilter<TInput,TOutput>
::PrecomputeGaussian()
{
	// Allocate memory
	if ( m_PreComputedGaussianKernelSize != m_KernelSize )
	{
		// Create new container and update the parameters
		m_PreComputedGaussianContainer = PreComputedGaussianContainerType::New();
		PreComputedGaussianContainerType::SizeType Size;
		Size[0] = 2*m_KernelSize+1;
		Size[1] = 2*m_KernelSize+1;
		m_PreComputedGaussianContainer->SetContainerSize(Size);
		m_PreComputedGaussianContainer->Reserve(Size[0]*Size[1]);
		m_PreComputedGaussianKernelSize = m_KernelSize;
		m_PreComputedGaussianSigma = 0;
	}

	// Precompute the gaussian
	if ( abs(m_PreComputedGaussianSigma-m_SpatialSigma) > 1e-6 )
	{
		m_PreComputedGaussianSigma = m_SpatialSigma;

		// Compute
		float *tmp = m_PreComputedGaussianContainer->GetImportPointer();
		for ( int i = 0; i < m_PreComputedGaussianKernelSize*2+1; i++ )
		{
			for ( int j = 0; j < m_PreComputedGaussianKernelSize*2+1; j++ )
			{
				double f1 = 1.f / 2.f*3.1415f*m_PreComputedGaussianSigma*m_PreComputedGaussianSigma;
				double f2 = i-m_PreComputedGaussianKernelSize;
				double f3 = j-m_PreComputedGaussianKernelSize;
				double f4 = sqrt(f2*f2+f3*f3);
				double f5 = expf(-0.5f*f4/(m_PreComputedGaussianSigma*m_PreComputedGaussianSigma));
				double w_d = f1*f5;

				tmp[i*(m_PreComputedGaussianKernelSize*2+1)+j] = w_d;
			}
		}

		// Transfer to device
		m_PreComputedGaussianContainer->HostModified();
		m_PreComputedGaussianContainer->Synchronize();
	}
}


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

	// Allocate CUDA device memory
	if ( m_CurrentSize != Size )
	{
		// Create new input image container
		m_InputImageContainer = InputImageContainerType::New();
		m_InputImageContainer->SetContainerSize(Size);
		m_InputImageContainer->Reserve(Size[0]*Size[1]);

		// Create output container
		m_OutputImageContainer = OutputImageContainerType::New();
		m_OutputImageContainer->SetContainerSize(Size);
		m_OutputImageContainer->Reserve(Size[0]*Size[1]);
		this->GetOutput()->SetPixelContainer(m_OutputImageContainer);

		// Update the size member
		m_CurrentSize = Size;
	}	
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
void
CUDABilateralImageFilter<TInput,TOutput>
::DoFiltering()
{
	// Precompute the gaussian
	PrecomputeGaussian();

	// Copy the data
	m_InputImageContainer->DeepCopy(this->GetInput()->GetPixelContainer());

	// Kernel invocation
	BilateralFilter
	(
		m_InputImageContainer->GetCudaMemoryArray(),
		m_PreComputedGaussianContainer->GetCudaMemoryArray(),
		m_OutputImageContainer->GetCudaMemoryPointer(),
		m_OutputImageContainer->GetPitch(),
		this->GetInput()->GetBufferedRegion().GetSize()[0],
		this->GetInput()->GetBufferedRegion().GetSize()[1],
		m_KernelSize,
		m_SpatialSigma,
		m_RangeSigma
	);

	// Due to fps evaluation
	//cudaThreadSynchronize();
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput>
void
CUDABilateralImageFilter<TInput,TOutput>
::GenerateData()
{
	// 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
	AllocateCudaMemory();

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









