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

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

#include "ritkDebugManager.h"

namespace ritk
{

//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>
::CudaPitchedMemoryImportImageContainer()
{
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>
::~CudaPitchedMemoryImportImageContainer()
{
	this->DeallocateManagedMemory();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
PrintSelf(std::ostream& os, itk::Indent indent) const
{
	os << "CudaPitchedMemoryImportImageContainer" << std::endl;
	os << indent << "Pitch=" << m_CudaMemory.front().second << std::endl;
	os << indent << "Memory @" << ((void*)m_CudaMemory.front().first) << std::endl;
	Superclass::PrintSelf(os,indent);
} 


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
typename CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::Element*
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
AllocateElements(ElementIdentifier size) const
{
	// Sanity checks. This ensures that we can operate on the internal 2-D size parameters
	if ( Superclass::m_ContainerSize[0] == 0 || Superclass::m_ContainerSize[1] == 0 )
	{
		LOG_DEB("Invalid width and/or height detected: did you forget to set them via SetSize()? RequestedSize="<<size);
		itkExceptionMacro(<<"Invalid width and/or height detected: did you forget to set them via SetSize()? RequestedSize="<<size);
	}
	if ( Superclass::m_ContainerSize[0]*Superclass::m_ContainerSize[1] != size )
	{
		LOG_DEB("Mismatch between widht/height and requested size detected: " << Superclass::m_ContainerSize[0] << "*" << Superclass::m_ContainerSize[1] << "=" << (Superclass::m_ContainerSize[0]*Superclass::m_ContainerSize[1]) << " vs. " << size);
		itkExceptionMacro(<<"Mismatch between widht/height and requested size detected: " << Superclass::m_ContainerSize[0] << "*" << Superclass::m_ContainerSize[1] << "=" << (Superclass::m_ContainerSize[0]*Superclass::m_ContainerSize[1]) << " vs. " << size);
	}

	// Execute in the global CUDA context. Note that we can't use the size argument directly. Instead, we use the internal 2-D size!
	CudaBlockedConstMethodCallMacro(AllocateCudaMemory);

	// Call the superclass' implementation. This also returns the correct host memory pointer
	return Superclass::AllocateElements(size);
}


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

	// Call the superclass' implementation to deallocate host memory
	Superclass::DeallocateManagedMemory();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
SynchronizeDevice()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(CopyHostToDevice);

	// Synchronize MTimes
	Superclass::SynchronizeDevice();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
SynchronizeHost()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(CopyDeviceToHost);

	// Synchronize MTimes
	Superclass::SynchronizeHost();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
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>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
DeepCopy(ImportImageContainerType *Container)
{
	// Call the superclass implementation first. This will ensure correct memory allocation and host memory copy.
	Superclass::DeepCopy(Container);

	// Check if we have a cuda container at all. If no, the GenericImportImageContainer should call a Synchronize() triggered by Superclass::DeepCopy() on us
	Superclass *CudaContainer = dynamic_cast<Superclass*>(Container);
	if ( !CudaContainer )
		return;

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

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


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
CopyToMappedResource(void *d_ptr, size_t size) const
{
	itkExceptionMacro(<<"Copying to mapped resources not supported by a pitched memory container!");
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
typename CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::Element*
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
GetCudaMemoryPointer()
{
	GenericImportImageContainerType::DeviceModified();
	return m_CudaMemory.front().first;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
const typename CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::Element*
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
GetCudaMemoryPointer() const
{
	return m_CudaMemory.front().first;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
size_t
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
GetPitch() const
{
	return m_CudaMemory.front().second;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
AllocateCudaMemory() const
{
	// Use our internal 2-D size for allocation. Its correctness is validated in AllocateElements()
	size_t Pitch;
	Element *Memory;
	ritkCudaSafeCall(cudaMallocPitch((void**)&Memory, &Pitch, Superclass::m_ContainerSize[0]*sizeof(Element), Superclass::m_ContainerSize[1]));

	// Save the allocated memory at the end of our internal list such that it won't be deleted through the itk (de)allocate scheme
	std::pair<Element*,size_t> tmp;
	tmp.first = Memory;
	tmp.second = Pitch;
	m_CudaMemory.push_back(tmp);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
DeallocateCudaMemory()
{
	// Sanity check
	if ( m_CudaMemory.size() == 0 )
	{
		return;
	}

	// Delete the first element in our internal list
	Element *Memory = m_CudaMemory.front().first;
	ritkCudaSafeCall(cudaFree((void*)Memory));
	m_CudaMemory.pop_front();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
CopyDeviceToHost()
{
	// Copy data from device to host
	ritkCudaSafeCall(cudaMemcpy2D(
		this->GetBufferPointer(),
		Superclass::m_ContainerSize[0]*sizeof(Element),
		this->GetCudaMemoryPointer(),
		this->GetPitch(),
		Superclass::m_ContainerSize[0]*sizeof(Element),
		Superclass::m_ContainerSize[1],
		cudaMemcpyDeviceToHost
	));
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
CopyHostToDevice()
{
	// Copy data from device to host
	ritkCudaSafeCall(cudaMemcpy2DAsync(
		this->GetCudaMemoryPointer(),
		this->GetPitch(),
		GenericImportImageContainerType::GetBufferPointer(),
		Superclass::m_ContainerSize[0]*sizeof(Element),
		Superclass::m_ContainerSize[0]*sizeof(Element),
		Superclass::m_ContainerSize[1],
		cudaMemcpyHostToDevice
	));
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaPitchedMemoryImportImageContainer<TElementIdentifier, TElement>::
DeepCopyCudaMemory()
{
	// Pitched -> Pitched
	const CudaPitchedMemoryImportImageContainer *PitchedContainer = dynamic_cast<const CudaPitchedMemoryImportImageContainer*>(m_TmpCudaContainer);
	if ( PitchedContainer )
	{
		// Sanity checks
		if ( PitchedContainer->m_ContainerSize[0] != Superclass::m_ContainerSize[0] || PitchedContainer->m_ContainerSize[1] != Superclass::m_ContainerSize[1] )
		{
			LOG_DEB("Invalid width and/or height detected: [" << PitchedContainer->m_ContainerSize[0] << "x" << PitchedContainer->m_ContainerSize[1] << "] vs. [" << Superclass::m_ContainerSize[0]  << "x" << Superclass::m_ContainerSize[1]  << "]");
			itkExceptionMacro(<<"Invalid width and/or height detected: [" << PitchedContainer->m_ContainerSize[0] << "x" << PitchedContainer->m_ContainerSize[1] << "] vs. [" << Superclass::m_ContainerSize[0]  << "x" << Superclass::m_ContainerSize[1]  << "]");
		}
	
		// Get pitch of the input container
		size_t Pitch = PitchedContainer->GetPitch();

		// Copy
		ritkCudaSafeCall(cudaMemcpy2DAsync
		(
			this->GetCudaMemoryPointer(),
			this->GetPitch(),
			PitchedContainer->GetCudaMemoryPointer(),
			Pitch,
			Superclass::m_ContainerSize[0]*sizeof(Element),
			Superclass::m_ContainerSize[1],
			cudaMemcpyDeviceToDevice
		));

		// Update the timestamps. Take care if host and device are already synchronized!
		GenericImportImageContainerType::DeviceModified();
		if ( m_TmpCudaContainer->GetDeviceTimeStamp() == m_TmpCudaContainer->GetHostTimeStamp() )
			GenericImportImageContainerType::m_HostMTime = GenericImportImageContainerType::m_DeviceMTime;

		// Clean up
		m_TmpCudaContainer = NULL;
		return;
	}

	// Regular -> Pitched
	typedef CudaRegularMemoryImportImageContainer<ElementIdentifier, Element> CudaRegularMemoryImportImageContainerType;
	const CudaRegularMemoryImportImageContainerType *RegularContainer = dynamic_cast<const CudaRegularMemoryImportImageContainerType*>(m_TmpCudaContainer);
	if ( RegularContainer )
	{
		// Copy
		ritkCudaSafeCall(cudaMemcpy2DAsync(
			this->GetCudaMemoryPointer(),
			this->GetPitch(),
			RegularContainer->GetCudaMemoryPointer(),
			Superclass::m_ContainerSize[0]*sizeof(Element),
			Superclass::m_ContainerSize[0]*sizeof(Element),
			Superclass::m_ContainerSize[1],
			cudaMemcpyDeviceToDevice
		));

		// Update the timestamps. Take care if host and device are already synchronized!
		GenericImportImageContainerType::DeviceModified();
		if ( m_TmpCudaContainer->GetDeviceTimeStamp() == m_TmpCudaContainer->GetHostTimeStamp() )
			GenericImportImageContainerType::m_HostMTime = GenericImportImageContainerType::m_DeviceMTime;

		// Clean up
		m_TmpCudaContainer = NULL;
		return;
	}

	// Array -> Pitched
	typedef CudaArrayImportImageContainer<ElementIdentifier, Element> CudaArrayImportImageContainerType;
	const CudaArrayImportImageContainerType *ArrayContainer = dynamic_cast<const CudaArrayImportImageContainerType*>(m_TmpCudaContainer);
	if ( ArrayContainer )
	{
		// Sanity check
		if ( ArrayContainer->GetContainerSize()[0] != Superclass::m_ContainerSize[0] || ArrayContainer->GetContainerSize()[1] != Superclass::m_ContainerSize[1] )
		{
			LOG_DEB("Invalid width and/or height detected: [" << ArrayContainer->GetContainerSize()[0] << "x" << ArrayContainer->GetContainerSize()[1] << "] vs. [" << Superclass::m_ContainerSize[0]  << "x" << Superclass::m_ContainerSize[1]  << "]");
			itkExceptionMacro(<<"Invalid width and/or height detected: [" << ArrayContainer->GetContainerSize()[0] << "x" << ArrayContainer->GetContainerSize()[1] << "] vs. [" << Superclass::m_ContainerSize[0]  << "x" << Superclass::m_ContainerSize[1]  << "]");
		}

		// Copy
		ritkCudaSafeCall(cudaMemcpy2DFromArrayAsync
		(
			this->GetCudaMemoryPointer(),
			this->GetPitch(),
			ArrayContainer->GetCudaMemoryArray(),
			0,
			0,
			Superclass::m_ContainerSize[0]*sizeof(Element),
			Superclass::m_ContainerSize[1],
			cudaMemcpyDeviceToDevice
		));

		// Update the timestamps. Take care if host and device are already synchronized!
		GenericImportImageContainerType::DeviceModified();
		if ( m_TmpCudaContainer->GetDeviceTimeStamp() == m_TmpCudaContainer->GetHostTimeStamp() )
			GenericImportImageContainerType::m_HostMTime = GenericImportImageContainerType::m_DeviceMTime;

		// Clean up
		m_TmpCudaContainer = NULL;
		return;
	}

	// We can't handle this container
	LOG_DEB("Invalid container detected: " << m_TmpCudaContainer);
	itkExceptionMacro(<<"Invalid container detected: " << m_TmpCudaContainer);
}


} // ritk
