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

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

#include "ritkCudaArrayImportImageContainer.h"
#include "ritkCudaRegularMemoryImportImageContainer.h"
#include "ritkCudaPitchedMemoryImportImageContainer.h"

#include "ritkDebugManager.h"

namespace ritk
{

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


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


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
PrintSelf(std::ostream& os, itk::Indent indent) const
{
	os << "CudaArrayImportImageContainer" << std::endl;
	os << indent << "Array @" << m_CudaArray.front().first << std::endl;
	Superclass::PrintSelf(os,indent);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
const cudaChannelFormatDesc&
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
GetChannelDesc() const
{
	return m_CudaArray.front().second;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
cudaChannelFormatDesc&
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
GetChannelDesc() 
{
	return m_CudaArray.front().second;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
typename CudaArrayImportImageContainer<TElementIdentifier, TElement>::Element*
CudaArrayImportImageContainer<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()?");
		itkExceptionMacro(<<"Invalid width and/or height detected: did you forget to set them via SetSize()?");
	}
	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
CudaArrayImportImageContainer<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
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
SynchronizeDevice()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(CopyHostToDevice);

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


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

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


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaArrayImportImageContainer<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
CudaArrayImportImageContainer<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. Eventually, 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
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
CopyToMappedResource(void *d_ptr, size_t size) const
{
	itkExceptionMacro(<<"Copying to mapped resources not supported by an array container!");
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
cudaArray*
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
GetCudaMemoryArray()
{
	GenericImportImageContainerType::DeviceModified();
	return m_CudaArray.front().first;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
const cudaArray*
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
GetCudaMemoryArray() const
{
	return m_CudaArray.front().first;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
AllocateCudaMemory() const
{
	// Create the channel descriptor for the TElement argument.
	cudaChannelFormatDesc Desc = cudaCreateChannelDesc<TElement>();

	// Allocate. Note that a correct size is checked in AllocateElements()
	cudaArray* Array;
	ritkCudaSafeCall(cudaMallocArray((cudaArray**)&Array, &Desc, Superclass::m_ContainerSize[0], 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<cudaArray*, cudaChannelFormatDesc> tmp;
	tmp.first = Array;
	tmp.second = Desc;
	m_CudaArray.push_back(tmp);
}


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

	// Delete the first element in our internal list
	cudaArray *Array = m_CudaArray.front().first;
	ritkCudaSafeCall(cudaFreeArray(Array));
	m_CudaArray.pop_front();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
CopyDeviceToHost()
{
	// Copy data from device to host
	long NumBytes = GenericImportImageContainerType::Capacity()*sizeof(Element);
	ritkCudaSafeCall(cudaMemcpyFromArray(GenericImportImageContainerType::GetBufferPointer(), this->GetCudaMemoryArray(), 0, 0, NumBytes, cudaMemcpyDeviceToHost));
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
CopyHostToDevice()
{
	// Copy data from host to device
	long NumBytes = GenericImportImageContainerType::Capacity()*sizeof(Element);
	ritkCudaSafeCall(cudaMemcpyToArrayAsync(this->GetCudaMemoryArray(), 0, 0, GenericImportImageContainerType::GetBufferPointer(), NumBytes, cudaMemcpyHostToDevice));
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
CudaArrayImportImageContainer<TElementIdentifier, TElement>::
DeepCopyCudaMemory()
{
	// Array -> Array
	const CudaArrayImportImageContainer *ArrayContainer = dynamic_cast<const CudaArrayImportImageContainer*>(m_TmpCudaContainer);
	if ( ArrayContainer )
	{
		// Copy
		ritkCudaSafeCall(cudaMemcpyArrayToArray
		(
			this->GetCudaMemoryArray(),
			0,
			0,
			ArrayContainer->GetCudaMemoryArray(),
			0,
			0,
			GenericImportImageContainerType::Capacity()*sizeof(Element),
			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 memory -> Array
	typedef CudaRegularMemoryImportImageContainer<ElementIdentifier, Element> CudaRegularMemoryImportImageContainerType;
	const CudaRegularMemoryImportImageContainerType *RegularContainer = dynamic_cast<const CudaRegularMemoryImportImageContainerType*>(m_TmpCudaContainer);
	if ( RegularContainer )
	{
		// Copy
		ritkCudaSafeCall(cudaMemcpyToArrayAsync
		(
			this->GetCudaMemoryArray(),
			0,
			0,
			RegularContainer->GetCudaMemoryPointer(),
			GenericImportImageContainerType::Capacity()*sizeof(Element),
			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;
	}

	// Pitched memory -> Array
	typedef CudaPitchedMemoryImportImageContainer<ElementIdentifier, Element> CudaPitchedMemoryImportImageContainerType;
	const CudaPitchedMemoryImportImageContainerType *PitchedContainer = dynamic_cast<const CudaPitchedMemoryImportImageContainerType*>(m_TmpCudaContainer);
	if ( PitchedContainer )
	{
		// Get pitch parameters
		size_t Pitch = PitchedContainer->GetPitch();
		size_t Width = PitchedContainer->GetContainerSize()[0];
		size_t Height = PitchedContainer->GetContainerSize()[1];

		// Copy
		ritkCudaSafeCall(cudaMemcpy2DToArrayAsync
		(
			this->GetCudaMemoryArray(),
			0,
			0,
			PitchedContainer->GetCudaMemoryPointer(),
			Pitch,
			Width*sizeof(Element),
			Height,
			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
