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

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkCudaImportImageContainer.txx
	Version:		$Revision: 2883 $
	Date:			$Date: 2012-04-02 17:01:22 +0200 (Mo, 02 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 "ritkCudaImportImageContainer.h"

#include "ritkDebugManager.h"

#include <cuda_gl_interop.h>

namespace ritk
{

//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>
::CudaImportImageContainer()
{
	// Make sure to initialize the CUDA context
	ritk::CudaContext::GetInstance();

	// Invalid memory by default
	m_ContainerSize.Fill(0);
	m_TmpCudaContainer = NULL;
	
	// No image set
	m_Image = NULL;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>
::~CudaImportImageContainer()
{
	this->DeallocateManagedMemory();

	// Unmap all resources
	for ( typename CUDAResourceContainerType::iterator it = m_CUDAResourceContainer.begin(); it != m_CUDAResourceContainer.end(); ++it )
	{
		// Get the resource
		CUDAResourceStruct &tmp = *it;
		
		// Unmap
		if ( tmp.Resource )
			cudaGraphicsUnregisterResource(tmp.Resource);
	}
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
PrintSelf(std::ostream& os, itk::Indent indent) const
{
	os << indent << "ContainerSize=[" << m_ContainerSize << std::endl;
	if ( m_HostMem.size() != 0 )
		os << indent << "HostMemory = " << ((void*)m_HostMem.front()) << std::endl;
	else
		os << indent << "No host memory..." << std::endl;
	Superclass::PrintSelf(os,indent);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
typename CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::Element*
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
GetBufferPointer()
{
	//cudaThreadSynchronize(); 
	return Superclass::GetBufferPointer();
}


#if 0
//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
const typename CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::Element*
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
GetBufferPointer() const
{
	//cudaThreadSynchronize(); 
	return Superclass::GetBufferPointer() const;
}
#endif


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
ReserveMemory(const SizeType &ContainerSize)
{
	this->SetContainerSize(ContainerSize);
	ElementIdentifier NumElems = 1;
	for ( unsigned int dim = 0; dim < Self::Dimension; dim++ )
		NumElems *= this->GetContainerSize()[dim];
	this->Reserve(NumElems);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
SetContainerSize(const SizeType &ContainerSize)
{
	m_ContainerSize = ContainerSize;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
typename CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::SizeType
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
GetContainerSize() const
{
	return m_ContainerSize;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
SetImage(typename itk::Image<Element,VDimension>::ConstPointer Image)
{
	m_Image = Image;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
typename CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::Element*
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
AllocateElements(ElementIdentifier size) const
{
	// Sanity checks. This ensures that we can operate on the internal container size parameters
	if ( this->m_ContainerSize[0] == 0 )
	{
		LOG_DEB("Invalid size detected: did you forget to set them via SetSize()? RequestedSize="<<size);
		itkExceptionMacro(<<"Invalid size: did you forget to set them via SetSize()? RequestedSize="<<size);
	}
	ElementIdentifier Num = 1;
	for ( unsigned int dim = 0; dim < Self::Dimension; dim++ )
		Num *= this->m_ContainerSize[dim];
	if ( Num != size )
	{
		LOG_DEB("Mismatch between container size and requested size detected: " << Num << " vs. " << size);
		itkExceptionMacro(<<"Mismatch between widht/height and requested size detected: " << Num << " vs. " << size);
	}

	// Execute in the global CUDA context
	CudaBlockedConstMethodCallMacro(AllocateCudaMemory);

	// Return the memory pointer that was pushed back by the previous call
	return m_HostMem.back();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
DeallocateManagedMemory()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(DeallocateCudaMemory);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
SynchronizeDevice()
{
	// Synchronize MTimes
	Superclass::SynchronizeDevice();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
SynchronizeHost()
{
	// Synchronize MTimes
	Superclass::SynchronizeHost();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
DeepCopy(const ImportImageContainerType *Container)
{
	// Cast and pass
	ImportImageContainerType *NonConstContainer = const_cast<ImportImageContainerType*>(Container);
	if ( !NonConstContainer )
	{
		LOG_DEB("Could not remove const qualifier!");
		itkExceptionMacro(<<"Could not remove const qualifier!");
	}
	DeepCopy(NonConstContainer);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
DeepCopy(ImportImageContainerType *Container)
{
	// Copying is depending on the actual type of the input container
	Self *CudaContainer = dynamic_cast<Self*>(Container);
	if ( !CudaContainer )
	{
		Superclass::DeepCopy(Container);
		return;
	}

	// Allocate memory is mandatory
	if ( this->Capacity() != Container->Capacity() || m_ContainerSize != CudaContainer->GetContainerSize() )
	{
		DeallocateManagedMemory();
		SetContainerSize(CudaContainer->GetContainerSize());
		Superclass::Reserve(Container->Capacity());
	}

	// Check for MTimes: we do not have to copy host memory that is out-of-date. Eventually, this has to be resolved by a synch!
	if ( CudaContainer->GetHostMTime() < CudaContainer->GetDeviceMTime() )
	{
		GenericImportImageContainerType::m_HostMTime = CudaContainer->GetHostTimeStamp();
		return;
	}

	// Host copy
	m_TmpCudaContainer = CudaContainer;
	CudaBlockedMethodCallMacro(DeepCopyCudaMemory);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
RegisterVBO(GLuint target) const
{
	// Sanity check if we already registered this VBO
	for ( typename CUDAResourceContainerType::iterator it = m_CUDAResourceContainer.begin(); it != m_CUDAResourceContainer.end(); ++it )
	{
		// Get the resource
		CUDAResourceStruct &tmp = *it;
		
		// Actual check
		if ( tmp.VBO == target )
		{
			LOG_DEB("VBO #" << target << " already registered...");
			itkExceptionMacro(<<"VBO #" << target << " already registered...");
			return;
		}
	}

	// Register this VBO
	CUDAResourceStruct tmp;
	tmp.VBO = target;
	ritkCudaSafeCall(cudaGraphicsGLRegisterBuffer(&tmp.Resource, tmp.VBO, cudaGraphicsMapFlagsWriteDiscard));

	// Save the resource
	m_CUDAResourceContainer.push_back(tmp);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
UnregisterVBO(GLuint target) const
{
	// Loop over the registered VBOs
	for ( typename CUDAResourceContainerType::iterator it = m_CUDAResourceContainer.begin(); it != m_CUDAResourceContainer.end(); ++it )
	{
		// Get the resource
		CUDAResourceStruct &tmp = *it;
		
		// Actual check
		if ( tmp.VBO == target )
		{
			ritkCudaSafeCall(cudaGraphicsUnregisterResource(tmp.Resource));
			m_CUDAResourceContainer.erase(it);
			return;
		}
	}

	// We don't have this VBO registered
	LOG_DEB("VBO #" << target << " not registered...");
	itkExceptionMacro(<<"VBO #" << target << " not registered...");
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
CopyToVBO(GLuint target, size_t size) const
{
	// Sanity check on the size
	if ( size != this->Capacity()*sizeof(Element) )
		itkExceptionMacro(<<"Invalid size detected: Capacity=" << this->Capacity()*sizeof(Element) << " RequestedSize=" << size);

	// Loop over the registered VBOs
	for ( typename CUDAResourceContainerType::iterator it = m_CUDAResourceContainer.begin(); it != m_CUDAResourceContainer.end(); ++it )
	{
		// Get the resource
		CUDAResourceStruct &tmp = *it;
		
		// Found the requested VBO
		if ( tmp.VBO == target )
		{
			// Map
			void *d_Ptr = NULL;
			ritkCudaSafeCall(cudaGraphicsMapResources(1, &tmp.Resource));
			ritkCudaSafeCall(cudaGraphicsResourceGetMappedPointer(&d_Ptr, &size, tmp.Resource));

			// Call the subclass' implementation
			CopyToMappedResource(d_Ptr, size);

			// Unmap
			ritkCudaSafeCall(cudaGraphicsUnmapResources(1, &tmp.Resource));

			// Finished here
			return;
		}
	}

	// We don't have this VBO registered
	LOG_DEB("VBO #" << target << " not registered...");
	std::cerr << "VBO #" << target << " not registered..." << std::endl;
	itkExceptionMacro(<<"VBO #" << target << " not registered...");
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
AllocateCudaMemory() const
{
	// Allocate. Note that a correct size is checked in AllocateElements()
	ElementIdentifier Num = 1;
	for ( unsigned int dim = 0; dim < Self::Dimension; dim++ )
		Num *= this->m_ContainerSize[dim];
	Element *Mem;
	ritkCudaSafeCall(cudaMallocHost((void**)&Mem, sizeof(Element)*Num));

	// Save the allocated memory at the end of our internal list such that it won't be deleted through the itk (de)allocate scheme
	m_HostMem.push_back(Mem);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
DeallocateCudaMemory()
{
	// Encapsulate all image memory deallocation here
	if ( m_HostMem.size() != 0 && this->GetContainerManageMemory() )
	{
		ritkCudaSafeCall(cudaFreeHost(m_HostMem.front()));
		m_HostMem.pop_front();
	}
	
	// Invalidate the import pointer and capacity. Note that we don't invalidate the container size on purpose.
	this->SetImportPointer(0);
	this->SetCapacity(0);
	this->SetSize(0);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement, unsigned int VDimension>
void
CudaImportImageContainer<TElementIdentifier,TElement,VDimension>::
DeepCopyCudaMemory()
{
	// Simple copy
	ritkCudaSafeCall(cudaMemcpy
	(
		this->GetBufferPointer(),
		m_TmpCudaContainer->GetBufferPointer(),
		this->Capacity()*sizeof(Element),
		cudaMemcpyHostToHost
	));

	// Clean up, and mark the host as modified
	GenericImportImageContainerType::HostModified();
	m_TmpCudaContainer = NULL;
}

} // ritk







