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

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkCudaImportImageContainer.h
	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

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

#ifndef CUDAIMPORTIMAGECONTAINER_H__
#define CUDAIMPORTIMAGECONTAINER_H__


#include "ritkCuda.h"
#include "ritkCudaContext.h"
#include "ritkGraphicsInteropImportImageContainer.h"

#include <list>

#include <itkImage.h>


namespace ritk
{

/**	@ingroup	Cuda
 *	@class		CudaImportImageContainer
 *	@author		Jakob Wasza
 *	@brief		Class extending the GenericImportImageContainer by CUDA GPU memory
 *
 *	@details
 *	This class provides basic interfaces for CUDA GPU memory management in compliance with the GenericImportImageContainer concept.  \n
 *	As there is no "single" CUDA memory type (there exists regular, pitched and array memory) and different dimensions have to be handled individually
 *	this class does not provide any direct methods for allocating/deallocating CUDA memory or synchronization. This is the task of subclasses.
 *	However, host memory allocation is performed by this base class. \n
 *	Memory (de)allocation in ITK as implemented in the Reserve() method is rather cumbersome: 
 *	as the first step, the new memory is obtained by calling AllocateElements() and old data copied to this new memory. Note that calling this
 *	method does not modify the internal memory pointer but returns a pointer to the allocated memory. As the second step, DeallocateManagedMemory() is called
 *	to release any previously allocated memory. However, in contrast to AllocateElements(), the DeallocateManagedMemory() method operates on the internal memory pointer.
 *	As the third and last step, the memory pointer obtained from AllocateElements() is set to the internal memory pointer. \n
 *	Therefore, host memory in this class is stored in a list. Allocated memory will always be pushed to the back whereas deallocating will always be performed
 *	on the list's front. Note that any sublass should use the same techique!
 *	Note that future versions of this class might overload the Reserve() method directly!\n
 *	An imporant aspect in memory allcoation with CUDA concerns the image size. As for certain CUDA memory types such as pitched memory or arrays the "flat" size of the
 *	data container is not sufficient but exact knowledge of the n-D size of the container is required, this base class provides the SetContainerSize() method that
 *	handles this task.
 *	\code
 * typedef CudaImportImageContainer<float,2> ContainerType;
 * ContainerType::Pointer Container = ContainerType::New();
 * Container->SetContainerSize(Image->GetRequestedRegion()); // this will set an internal member variable to be used by subclasses to determine correct memory allocation
 * Container->Reserve(Image->GetRequestedRegion().GetNumberOfPixels()); // standard ITK memory allocation interface but is delegated to subclasses
 *	\endcode
 *	For convenience this class provides a method that wraps these two steps via one Reserve(itk::Size<VDimension>) call. \n
 *	Note that due to performance reasons device to host (and vice versa) transfers may be implemented in subclasses using asynchronous calls.
 *	This does not pose any problem for asynchronous kernel invocations (and probably asynchronous device/device transfers), 
 *	but will probably result in invalid data states (host and device) if the host memory is read/written before the transfer is completed.
 *	This base class therefore overloads the GetBufferPointer() method ensuring synchronization. Note that this is the only "safe" way to access the internal CPU data.
 *	@sa SetContainerSize()
 *	@sa Reserve()
 *	@todo	Check for latest ITK version to support a more appropriate memory management (previously missing virtual qualifiers)
 *	@todo	Better async transfer support (e.g. streams)
 *	@todo	Check for sequential device to device asyn calls
 */
template<typename TElementIdentifier, typename TElement, unsigned int VDimension=2>
class CudaImportImageContainer : public GraphicsInteropImportImageContainer<TElementIdentifier,TElement>
{
public:
	/**	@name	Std typedefs */
	//@{
	typedef CudaImportImageContainer											Self;
	typedef GraphicsInteropImportImageContainer<TElementIdentifier,TElement>	Superclass;
	typedef itk::SmartPointer<Self>												Pointer;
	typedef itk::SmartPointer<const Self>										ConstPointer;
	//@}

	/// For convenience
	//@{
	typedef typename Superclass::GenericImportImageContainerType				GenericImportImageContainerType;
	typedef typename Superclass::ImportImageContainerType						ImportImageContainerType;
	//@}

	/**	@name	Template and size parameters */
	//@{
	typedef TElementIdentifier													ElementIdentifier;
	typedef TElement															Element;
	typedef itk::Size<VDimension>												SizeType;
	//@}

	/** @name Dimension macros */
	//@{
	itkStaticConstMacro(Dimension, unsigned int, VDimension);
	//@}

	/// RTTI
	itkTypeMacro(CudaImportImageContainer, GenericImportImageContainer);

	/**	@name Overloaded in order to support synchronization of asyn memory transfer calls */
	//@{
	virtual Element *GetBufferPointer();
	//virtual const Element *GetBufferPointer() const;
	//@}


	/**	@brief	Reserve memory
	 *	@param	Size	the n-D size of the container
	 *
	 *	@details
	 *	This is a convenience method that wraps calling SetContainerSize() followed by Reserve().
	 *	@todo	Compiler error when the method name is Reserve(). Maybe a namespace problem
	 */
	virtual void ReserveMemory(const SizeType &ContainerSize);

	/**	@name Set/get the n-D size (in pixel) of the container. This is necessary for the different CUDA memory types (i.e. pitched memory, CUDA arrays)*/
	//@{
	virtual void SetContainerSize(const SizeType &ContainerSize);
	virtual SizeType GetContainerSize() const;
	//@}

	/**	@brief	Set the containers image 
	 *	@param	Image	the image
	 *
	 *	@details
	 *	Set the image of this container. This method serves two purposes: \n
	 *		-	obviate the need to set the container size manually via. SetContainerSize() \n
	 *		-	workaround for allocating images as the method is unfortuantely not virtual in itk \n
	 *	For the first scenario, the image is set and the container size derived from the image's regions.
	 *	In the second scenario, the given parameter serves a the "owner" of the container. \n
	 *	For any Reserve() calls to this container, the container size can thus be determined by the owner's region. \n
	 *	Note that the container size is NOT set upon calling this method but on calling the reserve method! \n
	 *	IMPORTANT: this method should be considered experimental (!!!) and used with care due to missing virtual declarations
	 *	for several itk methods. Do NOT use this method unless you're sure about it...
	 *	@sa	SetContainerSize()
	 *	@sa Reserve()
	 *	@sa	AllocateElements()
	 */
	virtual void SetImage(typename itk::Image<Element,VDimension>::ConstPointer Image);

	/**	@name	Host/device synchronization. Is delegated to the superclass */
	//@{
	virtual void SynchronizeDevice();
	virtual void SynchronizeHost();
	//@}

	/**	@name	Deep copy */
	//@{

	/**	@brief	Deep copy from a container
	 *	@param	Container	The container providing the source
	 *
	 *	@details
	 *	Copy data from an itk::ImportImageContainer, however, this class' implementation only addresses copying host memory.
	 *	If the argument can be casted to a CudaImportImageContainer copying is tailored to page-locked memory and copying will
	 *	performed if and only if the host time stamp is later or equal than the device time stamp. Copying device memory has to be provided
	 *	by subclasses. If the argument cannot be casted to a CudaImportImageContainer, this method will simply call the Superclass'
	 *	implementation. In either case, memory allocation will be performed. \n
	 *	Note that in case of equal timestamps subclasses have to set them properly as this is NOT supported in the baseclass' implementation.
	 */
	virtual void DeepCopy(ImportImageContainerType *Container);

	/**	@brief	Deep copy from a container
	 *	@param	Container	The container providing the source
	 *
	 *	@details
	 *	This method will simply remove the const qualifier (const_cast) of the input argument and pass it to the non-const DeepCopy(ImportImageContainerType *) method.
	 *	Note that a non-const input is necessary due to the synchronization mechanism. However, it is granted, that the memory (host or device) that is up-to-data
	 *	will NOT be modified!
	 */
	virtual void DeepCopy(const ImportImageContainerType *Container);

	//@}

	/**	@name	CUDA to VBO copying */
	//@{

	/**	@brief	Copy the CUDA data to the given VBO
	 *	@param	target	the target VBO
	 *	@param	size	the size in bytes
	 *
	 *	@details
	 *	Calling this method eventually copies data to the given VBO. However, as there are several memory types in CUDA, the actual copy operation is triggered
	 *	via the pure virtual CopyToMappedResource(void*) method. This base class therefore only accounts for registering the VBO as a CUDA graphics resource.
	 *	The workflow of a copy operation triggered by this method is:
	 *		(1)		Sanity check comparing the requested size and the allocated size
	 *		(2)		Register the given VBO as CUDA graphics resource and unregister a previously registered VBO
	 *		(3)		Map the graphics resource
	 *		(4)		Call the actual copy operation defined in subclasses
	 *		(5)		Unmap the graphics resource
	 *	Note that calling this method may entail a device synchronization if the device data is not up-to-date
	 *	@todo	more detaild documentation on the workflow (with RegisterVBO and UnregisterVBO)
	 */
	void CopyToVBO(GLuint target, size_t size) const;

	/**	@brief	Register the given VBO to this container.
	 *	@param	target	the target VBO
	 *
	 *	@details
	 *	Register the given VBO to this container. Note that only ONE container may have registered the VBO.
	 *	@todo	more detaild documentation on the workflow (with CopyToVBO and UnregisterVBO)
	 */
	void RegisterVBO(GLuint target) const;

	/**	@brief	Unregister the given VBO from this container.
	 *	@param	target	the target VBO
	 *
	 *	@details
	 *	Unregister the given VBO from this container. Note that only ONE container may have registered the VBO.
	 *	@todo	more detaild documentation on the workflow (with CopyToVBO and RegisterVBO)
	 */
	void UnregisterVBO(GLuint target) const;


	/**	@brief	The actual device memory to VBO copy operation
	 *	@param	d_ptr	the mapped device pointer to a VBO
	 *	@param	size	the size in bytes that can be written
	 *
	 *	@details
	 *	This method is called by the CopyToVBO method and has to perform the actual data copy operation. Note that the base class ensures that the image container's capacity
	 *	matches the requested size. Further note, that the size argument should be considered as a last check only as, depending on the actual CUDA memory type, a 1D size
	 *	may be useless!
	 */
	virtual void CopyToMappedResource(void *d_ptr, size_t size) const =0;

	//@}


protected:
	/// Constructor
	CudaImportImageContainer();
	/// Destructor
	virtual ~CudaImportImageContainer();

	/// Print
	virtual void PrintSelf(std::ostream& os, itk::Indent indent) const;


	/**	@brief	Allocate the buffer
	 *	@param	size	The size in elements to allocate
	 *
	 *	@details
	 *	This method will allocate the buffer. Calling this method allocates
	 *	the host memory using cudaHostAlloc() and not using an ordinary
	 *	new or malloc call. This is due to the page-locked performance gain
	 *	when transfering from host to device and vice versa. \n
	 *	Note that this method completely overrides any Superclass' implementation.
	 *	However, the host memory is accessible as usual. \n
	 *	Note that calling this method will NOT perform any device allocation: this is
	 *	the task of a subclass' implementation.
	 *	@sa DeallocateManagedMemory
	 */
	virtual Element* AllocateElements(ElementIdentifier size) const;

	/**	@brief	Free the buffer
	 *
	 *	@details
	 *	This method will free the buffer previously allocated using 
	 *	CudaImportImageContainer::AllocateElements(). In fact, calling this method
	 *	performs a call to cudaFreeHost() to  m_ImportPointer declared in itk::ImportImageContainer.
	 *	This is in compliance with the standard implementation as can be found in itk::ImportImageContainer. \n
	 *	Note that calling this method won't perform any device de-allocation at all: this is the task
	 *	of subclasses.
	 *	@sa AllocateElements
	 */
	virtual void DeallocateManagedMemory();

	/**	@name	Methods that have to be executed in the global CUDA context */
	//@{
	CudaConstMethodMacro(CudaImportImageContainer, AllocateCudaMemory);
	CudaMethodMacro(CudaImportImageContainer, DeallocateCudaMemory);
	Self *m_TmpCudaContainer;
	CudaMethodMacro(CudaImportImageContainer, DeepCopyCudaMemory);
	//@}

protected:
	/// The size in pixel of this container
	SizeType m_ContainerSize;

	/// The image for this container (experimental)
	typename itk::Image<Element, VDimension>::ConstPointer m_Image;

	/// List holding the memory. This is due to the (de)allocate scheme in itk. Always delete the front and push new memory
	mutable std::list<Element*> m_HostMem;

	/**	@name	CUDA to OpenGL members */
	//@{
	typedef struct
	{
		GLuint VBO;
		cudaGraphicsResource_t Resource;
	} CUDAResourceStruct;
	typedef std::list<CUDAResourceStruct> CUDAResourceContainerType;
	mutable CUDAResourceContainerType m_CUDAResourceContainer;
	//@}

private:
	/// purposely not implemented
	CudaImportImageContainer(const Self&); 
	/// purposely not implemented
	void operator=(const Self&);
};

/** @name Template instantiations */
//@{
typedef CudaImportImageContainer<unsigned long,float> CudaImportImageContainerF;
typedef CudaImportImageContainer<unsigned long, itk::Vector<float,3> > CudaImportImageContainerF3;
//@}

} // ritk


#endif // CUDAIMPORTIMAGECONTAINER_H__
