#include "KinectSDKImageIO.h"
#include "ritkDebugManager.h"

template<class TRImageBase, class TKinectSDKImage>
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
KinectSDKImageIO()
{
	m_CurrentHeader = 0;
	m_uiSkeletons = NUMBEROFSKELETONS;
	m_uiQualityFlags = 1;

	// debug
	std::cout << "KinectSDK IO object created" << std::endl;
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
~KinectSDKImageIO()
{
	if ( m_CurrentHeader )
		delete m_CurrentHeader;
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
const char* 
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
GetRImageClassName()
{
	return KinectSDKImageType::New()->GetNameOfClass();
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
unsigned int
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
GetNumberOfFrames()
{
	OpenFileStreamAndReadHeader();
	
	return m_CurrentHeader->NumFrames;
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
void
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
GetResolution(unsigned int *SX, unsigned int *SY)
{
	OpenFileStreamAndReadHeader();

	*SX = m_CurrentHeader->CurrCameraParameters[ritk::CameraParameters::SX];
	*SY = m_CurrentHeader->CurrCameraParameters[ritk::CameraParameters::SY];
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
std::string
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
GetFileName()
{
	return m_FileName;
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
void
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
SetFileName(std::string FileName)
{
	if ( FileName == m_FileName )
		return;

	// Update filename
	m_FileName = FileName;

	// Set dirty flags
	m_FileNameDirtyFlag = true;

	// Reset index
	m_Index = 0;
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
void
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
Close()
{
	if ( m_OutputStream.is_open() )
		m_OutputStream.close();
	if ( m_InputStream.is_open() )
		m_InputStream.close();
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
void
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
WriteFrame(RImageBaseConstPointer Frame)
{

	// Sanity checks
	if ( !Frame )
		itkExceptionMacro(<<"No input set!");
	if ( m_FileName.size() == 0 )
		itkExceptionMacro(<<"No file name set!");

	// Update the pipeline
	KinectSDKImagePointer NonConstInput = const_cast<KinectSDKImageType*>(dynamic_cast<const KinectSDKImageType*>((const RImageBaseType*)Frame));
	if ( !NonConstInput )
		itkExceptionMacro(<<"Invalid input detected: Is not a KinectSDK image!");

	// ???
	NonConstInput->Update();

	// Dirty flag => close old stream open new one and write the header
	if ( m_FileNameDirtyFlag )
	{
		// Close
		if ( m_OutputStream.is_open() )
			m_OutputStream.close();
		
		// Open new file
		m_OutputStream.open(GetFileName().c_str(),std::ios::out|std::ios::trunc|std::ios::binary);

		// Check for errors
		if ( !m_OutputStream.is_open() || m_OutputStream.fail() )
			itkExceptionMacro(<< "Could not open output file: <" << GetFileName() <<">");

		// Write the type
		m_OutputStream << GetRImageClassName() << std::endl;


		// Write the header
		m_OutputStream << "FormatVersion	" << "1.1" << std::endl;
		m_OutputStream << "CameraParameters	" << NonConstInput->GetCameraParameters() << std::endl;

		// The number of frames
		m_OutputStream << "NumFrames	";
		// Remember the position to the end of the header
		m_HeaderEndPos = m_OutputStream.tellp();
		//std::cout << pos = " << m_HeaderEndPos << std::endl;
		m_OutputStream << "000000" << std::endl;

		// Reset index
		m_Index = 0;

		// Finally, clear the dirty flag
		m_FileNameDirtyFlag = false;
	}

	// Update the number of frames at the correct position
	m_OutputStream.seekp(m_HeaderEndPos);
	m_Index++;
	if ( m_Index > 1000000 )
		itkExceptionMacro(<<"Can not handle so many frames...");
	if ( m_Index < 10 )
		m_OutputStream << "00000";
	else if ( m_Index < 100 )
		m_OutputStream << "0000";
	else if ( m_Index < 1000 )
		m_OutputStream << "000";
	else if ( m_Index < 10000 )
		m_OutputStream << "00";
	else if ( m_Index < 100000 )
		m_OutputStream << "0";
	m_OutputStream << m_Index << std::endl;

	// Write the data
	m_OutputStream.seekp(0,std::ios::end);
	unsigned int SizeX = NonConstInput->GetRequestedRegion().GetSize()[0];
	unsigned int SizeY = NonConstInput->GetRequestedRegion().GetSize()[1];
	// Time stamp
	m_OutputStream.write((char*)&NonConstInput->GetTimeStamp(),sizeof(RImageBaseType::TimeStampType)); 

	// Conversion parameters
	KinectSDKImageType::RealType XToZ;
	KinectSDKImageType::RealType YToZ;
	NonConstInput->GetConversionParameters(&XToZ,&YToZ);
	m_OutputStream.write((char*)&XToZ,sizeof(KinectSDKImageType::RealType));	
	m_OutputStream.write((char*)&YToZ,sizeof(KinectSDKImageType::RealType));

	// Superclass' data
	m_OutputStream.write((char*)NonConstInput->GetRGBImage()->GetBufferPointer(),SizeX*SizeY*sizeof(KinectSDKImageType::RGBImageType::PixelType));
	m_OutputStream.write((char*)NonConstInput->GetRangeImage()->GetBufferPointer(),SizeX*SizeY*sizeof(RImageBaseType::RangeImageType::PixelType));
	m_OutputStream.write((char*)NonConstInput->GetFlagImage()->GetBufferPointer(),SizeX*SizeY*sizeof(RImageBaseType::FlagImageType::PixelType));


	// Skeleton data (joint positions and CoM)
	for (int i = 0; i < m_uiSkeletons; i++)
	{
		for (Skeleton::JointTypes joint = Skeleton::SKEL_HIP_Center; joint != Skeleton::SKEL_COUNT; joint = static_cast<Skeleton::JointTypes> (joint + 1))
		{
			
			VF3 tmpJoint  = NonConstInput->GetSkeleton()->GetJoint(joint,i);
			short tmpTrackingState = NonConstInput->GetSkeleton()->GetTrackingState(joint,i);

			m_OutputStream.write((char *) &tmpJoint,sizeof(VF3));
			m_OutputStream.write((char *) &tmpTrackingState,sizeof(short));
		}
	}

	// Reposition stream
	m_OutputStream.seekp(0,std::ios::end);

	// Check for errors
	if ( m_OutputStream.fail() )
		itkExceptionMacro(<<"Failed during writing!");
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
typename KinectSDKImageIO<TRImageBase,TKinectSDKImage>::RImageBasePointer
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
ReadFrame(unsigned int FrameID)
{

	// Open the stream, read the header
	OpenFileStreamAndReadHeader();

	// Create the result
	KinectSDKImagePointer Result = KinectSDKImageType::New();

	// Set all regions and allocate
	RImageBaseType::SizeType ImageSize;
	ImageSize[0] = m_CurrentHeader->CurrCameraParameters[ritk::CameraParameters::SX];
	ImageSize[1] = m_CurrentHeader->CurrCameraParameters[ritk::CameraParameters::SY];
	RImageBaseType::IndexType Index;
	Index[0] = 0;
	Index[1] = 0;
	RImageBaseType::RegionType Region;
	Region.SetIndex(Index);
	Region.SetSize(ImageSize);
	Result->SetRegions(Region);
	Result->Allocate();

	// Set the parameters
	for ( unsigned int i = 0; i < ritk::CameraParameters::NumCameraParameters; i++ )
		Result->GetCameraParameters()[i] = m_CurrentHeader->CurrCameraParameters[i];

	// Sanity check
	if ( FrameID >= m_CurrentHeader->NumFrames )
		itkExceptionMacro(<<"Invalid frame requested: " << FrameID << " vs. " << m_CurrentHeader->NumFrames);

	// Safe the stream position
	std::ifstream::pos_type DataSectionStreamPos = m_InputStream.tellg();
	
	// Depending on the format version
	if ( m_CurrentHeader->FormatVersion == std::string("1.0") 
	 ||  m_CurrentHeader->FormatVersion == std::string("1.1") )
	{
		// Some sizes to calculate the offsets
		long NumPixels = ImageSize[0]*ImageSize[1];
		std::ifstream::pos_type TimeStampSize = sizeof(RImageBaseType::TimeStampType);
		std::ifstream::pos_type ConversionParamsSize = 2*sizeof(RImageBaseType::RealType);
		std::ifstream::pos_type RGBImageSize = NumPixels*sizeof(KinectSDKImageType::RGBType);
		std::ifstream::pos_type RangeImageSize = NumPixels*sizeof(RImageBaseType::RangeType);
		std::ifstream::pos_type FlagImageSize = NumPixels*sizeof(RImageBaseType::FlagType);
		std::ifstream::pos_type SkeletonDataSize = Skeleton::SKEL_COUNT * m_uiSkeletons * (sizeof(VF3) + m_uiQualityFlags * sizeof(short));

		// Full framesize
		std::ifstream::pos_type FrameSize = TimeStampSize + ConversionParamsSize + RGBImageSize + RangeImageSize + FlagImageSize + SkeletonDataSize;
		
		// Read the timestamp
		RImageBaseType::TimeStampType TimeStamp;
		m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize));
		m_InputStream.read((char*)&TimeStamp,TimeStampSize);
		//std::cout<< "Time: " << TimeStamp << std::endl;
		Result->SetTimeStamp(TimeStamp);

		// Conversion parameters
		KinectSDKImageType::RealType ConvParams[2];
		m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize));
		m_InputStream.read((char*)ConvParams,ConversionParamsSize);
		Result->SetConversionParameters(ConvParams[0],ConvParams[1]);

		// Read the RGB values
		m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize + ConversionParamsSize));
		m_InputStream.read((char*)Result->GetRGBImage()->GetBufferPointer(),RGBImageSize);
		Result->GetRGBImage()->Modified();

		// Read the range values
		m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize + ConversionParamsSize + RGBImageSize));
		m_InputStream.read((char*)Result->GetRangeImage()->GetBufferPointer(),RangeImageSize);
		Result->GetRangeImage()->Modified();

		// Read the flags
		m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize + ConversionParamsSize + RGBImageSize + RangeImageSize));
		m_InputStream.read((char*)Result->GetFlagImage()->GetBufferPointer(),FlagImageSize);
		Result->GetFlagImage()->Modified();

		std::ifstream::pos_type SkeletonBegin = DataSectionStreamPos + 
			(std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize + ConversionParamsSize + RGBImageSize + RangeImageSize + FlagImageSize);

		// Read the skeleton data
		for (int i = 0, j = 0; i < m_uiSkeletons; i++)
		{
			for (Skeleton::JointTypes joint = Skeleton::SKEL_HIP_Center; joint != Skeleton::SKEL_COUNT; joint = static_cast<Skeleton::JointTypes> (joint + 1), j++)
			{
				VF3 tmpJoint;
				short trackingState = 0;

				// go to joint position and read
				m_InputStream.seekg(SkeletonBegin + (std::ifstream::pos_type) (j * sizeof(VF3) + j * m_uiQualityFlags * sizeof(short)));
				m_InputStream.read((char *) &tmpJoint, sizeof(VF3));

				
				if (m_uiQualityFlags)
				{
					// go to tracking State position and read
					m_InputStream.seekg(SkeletonBegin + (std::ifstream::pos_type) ((j+1) * sizeof(VF3) + j * sizeof(short)));
					m_InputStream.read((char *) &trackingState, sizeof(short));

				}
				
				// set data
				Result->GetSkeleton()->SetJoint(joint, tmpJoint, i);
				Result->GetSkeleton()->SetTrackingState(joint, trackingState, i);
				
			}
		}

		// go to position
		m_InputStream.seekg(DataSectionStreamPos + 
				(std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize + ConversionParamsSize + RGBImageSize + RangeImageSize + FlagImageSize + SkeletonDataSize));
				
		Result->GetSkeleton()->Modified();

		// Restore the stream position
		m_InputStream.seekg(DataSectionStreamPos);
	}
	else
	{
		itkExceptionMacro(<<"Unknown KinectSDKImageIO version: " << m_CurrentHeader->FormatVersion);
	}

	// Check for errors
	if ( m_InputStream.fail() )
		itkExceptionMacro(<<"Could not read frame!");

	return dynamic_cast<RImageBaseType*>((KinectSDKImageType*)Result);
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
void
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
OpenFileStreamAndReadHeader()
{
	// If the dirty flag for the filename is not set we don't have to do anything
	if ( !m_FileNameDirtyFlag )
		return;

	// Sanity check
	if ( m_FileName.size() == 0 )
		itkExceptionMacro(<<"No file name set!");

	// Close old stream and open new stream
	if ( m_InputStream.is_open() )
		m_InputStream.close();
	m_InputStream.open(GetFileName().c_str(),std::ios::in|std::ios::binary);
	if ( !m_InputStream.is_open() )
		itkExceptionMacro(<<"Could not open input file!");

	// Read the header
	if ( m_CurrentHeader )
		delete m_CurrentHeader;
	m_CurrentHeader = new HeaderStruct();
	std::string dummy;

	// Image type
	m_InputStream >> dummy;
	if ( dummy.compare(std::string(GetRImageClassName())) != 0 )
		itkExceptionMacro(<<"Invalid image type detected: " << dummy << " Can only handle " << GetRImageClassName() << "!");

	// Format version
	m_InputStream >> dummy;
	m_InputStream >> m_CurrentHeader->FormatVersion;
	LOG_DEB("File <" << m_FileName << ">: KinectSDKImageIO version " << m_CurrentHeader->FormatVersion);

	// Based on the old format version
	if ( m_CurrentHeader->FormatVersion == std::string("1.0") )
	{
		m_uiSkeletons = 1;
		m_uiQualityFlags = 0;
	}
	else if ( m_CurrentHeader->FormatVersion == std::string("1.1") )
	{
	}
	else
	{
		itkExceptionMacro(<<"Unknown KinectSDKImageIO version: " << m_CurrentHeader->FormatVersion);
	}

	// The camera parameters
	m_InputStream >> dummy;
	m_InputStream >> m_CurrentHeader->CurrCameraParameters;

	// The number of available frames
	m_InputStream >> dummy;
	m_InputStream >> m_CurrentHeader->NumFrames;

	// Set the stream position to the beginning of the data section
	m_InputStream.seekg(((unsigned long)m_InputStream.tellg())+1);

	// Finally, clear the dirty flag
	m_FileNameDirtyFlag = false;

	// Initializing the timestamps is done here as well
	InitTimeStamps();

	// Check for errors
	if ( m_InputStream.fail() )
		itkExceptionMacro(<<"Could not read header");
}


//----------------------------------------------------------------------------
template<class TRImageBase, class TKinectSDKImage>
void
KinectSDKImageIO<TRImageBase,TKinectSDKImage>::
InitTimeStamps()
{
	// debug
	std::cout << "Initializing timestamps" << std::endl;

	LOG_DEB("Initializing timestamps...");
	// Safe the stream position
	std::ifstream::pos_type DataSectionStreamPos = m_InputStream.tellg();

	// Depending on the current header
	if (m_CurrentHeader->FormatVersion == std::string("1.0")  || 
		m_CurrentHeader->FormatVersion == std::string("1.1"))
	{
		// Some offsets
		long NumPixels = m_CurrentHeader->CurrCameraParameters[ritk::CameraParameters::SX]*m_CurrentHeader->CurrCameraParameters[ritk::CameraParameters::SY];
		std::ifstream::pos_type TimeStampSize = sizeof(RImageBaseType::TimeStampType);
		std::ifstream::pos_type ConversionParamsSize = 2*sizeof(RImageBaseType::RealType);
		std::ifstream::pos_type RGBImageSize = NumPixels*sizeof(KinectSDKImageType::RGBType);
		std::ifstream::pos_type RangeImageSize = NumPixels*sizeof(RImageBaseType::RangeType);
		std::ifstream::pos_type FlagImageSize = NumPixels*sizeof(RImageBaseType::FlagType);
		std::ifstream::pos_type SkeletonDataSize = Skeleton::SKEL_COUNT * m_uiSkeletons * (sizeof(VF3) + m_uiQualityFlags * sizeof(short));

		std::ifstream::pos_type FrameSize = TimeStampSize + ConversionParamsSize + RGBImageSize + RangeImageSize + FlagImageSize + SkeletonDataSize;

		// Extract timestamps
		m_TimeStamps.resize(m_CurrentHeader->NumFrames);
		for ( long FrameID = 0; FrameID < m_CurrentHeader->NumFrames; FrameID++ )
		{
			RImageBaseType::TimeStampType TimeStamp;
			m_InputStream.seekg(DataSectionStreamPos + FrameID*FrameSize);
			m_InputStream.read((char*)&TimeStamp,TimeStampSize);
			m_TimeStamps[FrameID] = TimeStamp;
		}
	}
	else
	{
		itkExceptionMacro(<<"Unknown KinectSDKImageIO version: " << m_CurrentHeader->FormatVersion);
	}

	// Restore the stream position
	m_InputStream.seekg(DataSectionStreamPos);
}