#include "CUDATemporalAveragingImageFilter.h"

#include "ritkCudaRegularMemoryImportImageContainer.h"


// CUDA interface for element-wise addition: Accu[i] += Data[i]
//----------------------------------------------------------------------------
void AddData(float *Accu, float *Data, unsigned long Num);
void AddData(itk::Vector<float,3> *Accu, itk::Vector<unsigned char,3> *Data, unsigned long Num);


// CUDA interface for element-wise subtraction: Accu[i] -= Data[i]
//----------------------------------------------------------------------------
void SubData(float *Accu, float *Data, unsigned long Num);
void SubData(itk::Vector<float,3> *Accu, itk::Vector<unsigned char,3> *Data, unsigned long Num);


// CUDA interface for scalar multiplication: Output[i] = Scalar*Input[i]
//----------------------------------------------------------------------------
void ScalarMulData(float *Output, float *Input, float Scalar, unsigned long Num);
void ScalarMulData(itk::Vector<unsigned char,3> *Output, itk::Vector<float,3> *Input, float Scalar, unsigned long Num);


// CUDA interface for the recursive implementation
//----------------------------------------------------------------------------
void SlidingWindowTemporalAveraging(float *Output, float *Accu, float *AddData, float *SubData, float Scale, unsigned long Num);
void SlidingWindowTemporalAveraging(itk::Vector<unsigned char,3> *Output, itk::Vector<float,3> *Accu, itk::Vector<unsigned char,3> *AddData, itk::Vector<unsigned char,3> *SubData, float Scale, unsigned long Num);




//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::CUDATemporalAveragingImageFilter()
{
	m_NumFrames = 10;

	// Cuda memory
	m_CurrentSize.Fill(0);
	m_MemoryPool = NULL;
	m_AccuContainer = NULL;
	m_OutputContainer = NULL;
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::~CUDATemporalAveragingImageFilter()
{
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
void
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::Reset()
{
	// Clear queue
	m_FrameQueue = std::queue<unsigned int>();
	m_CurrentSize.Fill(0);
	m_MemoryPool = NULL;

	this->Modified();
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
void
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::SetNumFrames(unsigned int NumFrames)
{
	m_NumFrames = NumFrames;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
unsigned int
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::GetNumFrames()
{
	return m_NumFrames;
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
void
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::AllocateMemory()
{
	// Reset the filter. This will also free previously allocated memory
	Reset();

	// Initialize the size member
	m_CurrentSize = this->GetInput()->GetBufferedRegion().GetSize();

	// Initialize accu and mean image
	typename AccuImageContainerType::SizeType AccuContainerSize;
	AccuContainerSize[0] = m_CurrentSize[0]*m_CurrentSize[1];
	m_AccuContainer = AccuImageContainerType::New();
	m_AccuContainer->SetContainerSize(AccuContainerSize);
	m_AccuContainer->Reserve(m_CurrentSize[0]*m_CurrentSize[1]);
	ritkCudaSafeCall(cudaMemset(m_AccuContainer->GetCudaMemoryPointer(), 0, sizeof(typename AccuImageType::PixelType)*m_CurrentSize[0]*m_CurrentSize[1]));

	typename OutputImageContainerType::SizeType OutputContainerSize;
	OutputContainerSize[0] = m_CurrentSize[0]*m_CurrentSize[1];
	m_OutputContainer = OutputImageContainerType::New();
	m_OutputContainer->SetContainerSize(OutputContainerSize);
	m_OutputContainer->Reserve(m_CurrentSize[0]*m_CurrentSize[1]);
	ritkCudaSafeCall(cudaMemset(m_OutputContainer->GetCudaMemoryPointer(), 0, sizeof(typename InputImageType::PixelType)*m_CurrentSize[0]*m_CurrentSize[1]));
	this->GetOutput()->SetPixelContainer(m_OutputContainer);

	// Initialize memory pool
	m_MemoryPool = CUDAMemoryPoolType::New();
	m_MemoryPool->SetSlotSize(AccuContainerSize);
	m_MemoryPool->ResizePool(32);
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
void
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::DoFiltering()
{
	// Copy and enqueue the new frame (first we have to request a free slot)
	unsigned int ID = m_MemoryPool->RequestSlot();

	// As a next step, copy the data from the input
	m_MemoryPool->GetContainer(ID)->DeepCopy(this->GetInput()->GetPixelContainer());

	// Three cases: Queue underfill/overfill/filled
	if ( m_FrameQueue.size() < m_NumFrames )
	{
		// Add the new frame and normalize
		m_FrameQueue.push(ID);
		AddData(m_AccuContainer->GetCudaMemoryPointer(), m_MemoryPool->GetContainer(ID)->GetCudaMemoryPointer(), this->GetInput()->GetBufferedRegion().GetNumberOfPixels());
		ScalarMulData(m_OutputContainer->GetCudaMemoryPointer(), m_AccuContainer->GetCudaMemoryPointer(), 1.f/m_FrameQueue.size(), this->GetInput()->GetBufferedRegion().GetNumberOfPixels());
	}
	else if ( m_FrameQueue.size() > m_NumFrames )
	{
		// Successively remove the fronts
		while ( m_FrameQueue.size() > m_NumFrames )
		{
			// Remove from queue
			ID = m_FrameQueue.front();
			m_FrameQueue.pop();

			// Subtract (don't forget to release the memory)
			SubData(m_AccuContainer->GetCudaMemoryPointer(), m_MemoryPool->GetContainer(ID)->GetCudaMemoryPointer(), this->GetInput()->GetBufferedRegion().GetNumberOfPixels());
			m_MemoryPool->ReleaseSlot(ID);
		}

		// Normalize
		ScalarMulData(m_OutputContainer->GetCudaMemoryPointer(), m_AccuContainer->GetCudaMemoryPointer(), 1.f/m_FrameQueue.size(), this->GetInput()->GetBufferedRegion().GetNumberOfPixels());
	}
	else
	{
		// The sliding window
		m_FrameQueue.push(ID);
		SlidingWindowTemporalAveraging(m_OutputContainer->GetCudaMemoryPointer(), m_AccuContainer->GetCudaMemoryPointer(), m_MemoryPool->GetContainer(ID)->GetCudaMemoryPointer(), m_MemoryPool->GetContainer(m_FrameQueue.front())->GetCudaMemoryPointer(), 1.f/m_NumFrames, this->GetInput()->GetBufferedRegion().GetNumberOfPixels());
		m_MemoryPool->ReleaseSlot(m_FrameQueue.front());
		m_FrameQueue.pop();
	}
}


//----------------------------------------------------------------------------
template<class TInput,class TOutput,class TAccu>
void
CUDATemporalAveragingImageFilter<TInput,TOutput,TAccu>
::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
	if ( this->GetInput()->GetBufferedRegion().GetNumberOfPixels() != m_CurrentSize[0]*m_CurrentSize[1] )
		CudaBlockedMethodCallMacro(AllocateMemory);

	// Filtering
	CudaBlockedMethodCallMacro(DoFiltering);
}
