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

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkRImageIO.txx
	Version:		$Revision: 1761 $
	Date:			$Date: 2011-09-23 17:25:41 +0200 (Fr, 23 Sep 2011) $

	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 "ritkRImageIO.h"

#include <itkImageRegionConstIterator.h>

namespace ritk
{

template<class TRImageBase>
RImageIO<TRImageBase>::
RImageIO()
{
	m_CurrentHeader = 0;	
}


template<class TRImageBase>
RImageIO<TRImageBase>::
~RImageIO()
{
	if ( m_CurrentHeader )
		delete m_CurrentHeader;
}


//----------------------------------------------------------------------------
template<class TRImageBase>
const char* 
RImageIO<TRImageBase>::
GetRImageClassName()
{
	return RImageBaseType::New()->GetNameOfClass();
}


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


//----------------------------------------------------------------------------
template<class TRImageBase>
void
RImageIO<TRImageBase>::
GetResolution(unsigned int *SX, unsigned int *SY)
{
	OpenFileStreamAndReadHeader();
	
	*SX = (unsigned int)m_CurrentHeader->CurrCameraParameters[CameraParameters::SX];
	*SY = (unsigned int)m_CurrentHeader->CurrCameraParameters[CameraParameters::SY];
}


//----------------------------------------------------------------------------
template<class TRImageBase>
std::string
RImageIO<TRImageBase>::
GetFileName()
{
	return Superclass::m_FileName;
}


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

	// Update filename
	Superclass::m_FileName = FileName;

	// Set dirty flags
	m_FileNameDirtyFlag = true;

	// Reset index
	m_Index = 0;
}


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


//----------------------------------------------------------------------------
template<class TRImageBase>
void
RImageIO<TRImageBase>::
WriteFrame(RImageBaseConstPointer Frame)
{
	// Sanity checks
	if ( !Frame )
		itkExceptionMacro(<<"No input set!");
	if ( Superclass::m_FileName.size() == 0 )
		itkExceptionMacro(<<"No file name set!");

	// Update the pipeline
	RImageBasePointer NonConstInput = const_cast<RImageBaseType*>((const RImageBaseType*)Frame);
	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.0" << std::endl;
		m_OutputStream << "CameraParameters		" << NonConstInput->GetCameraParameters() << std::endl;
		m_OutputStream << std::endl;
		m_OutputStream << "NumFrames	";
		// Remember the position to the end of the header
		m_HeaderEndPos = m_OutputStream.tellp();
		//std::cout << "Header 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];
	m_OutputStream.write((char*)&NonConstInput->GetTimeStamp(),sizeof(typename RImageBaseType::TimeStampType));
	//m_OutputStream.write((char*)NonConstInput->GetRGBImage()->GetBufferPointer(),SizeX*SizeY*sizeof(RImageBaseType::RGBImageType::PixelType));
	m_OutputStream.write((char*)NonConstInput->GetRangeImage()->GetBufferPointer(),SizeX*SizeY*sizeof(typename RImageBaseType::RangeImageType::PixelType));
	m_OutputStream.write((char*)NonConstInput->GetFlagImage()->GetBufferPointer(),SizeX*SizeY*sizeof(typename RImageBaseType::FlagImageType::PixelType));
	m_OutputStream.seekp(0,std::ios::end);

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


//----------------------------------------------------------------------------
template<class TRImageBase>
typename RImageIO<TRImageBase>::RImageBasePointer
RImageIO<TRImageBase>::
ReadFrame(unsigned int FrameID)
{
	// Open the stream, read the header
	OpenFileStreamAndReadHeader();

	// Create the result
	RImageBasePointer Result = RImageBaseType::New();

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

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

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

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

	// Some offsets
	long NumPixels = ImageSize[0]*ImageSize[1];
	std::ifstream::pos_type TimeStampSize = sizeof(typename RImageBaseType::TimeStampType);
	//std::ifstream::pos_type RGBImageSize = NumPixels*sizeof(RImageBaseType::RGBType);
	std::ifstream::pos_type RangeImageSize = NumPixels*sizeof(typename RImageBaseType::RangeType);
	std::ifstream::pos_type FlagImageSize = NumPixels*sizeof(typename RImageBaseType::FlagType);
	//std::ifstream::pos_type FrameSize = TimeStampSize + RGBImageSize + RangeImageSize + FlagImageSize;
	std::ifstream::pos_type FrameSize = TimeStampSize + RangeImageSize + FlagImageSize;

	// Read the timestamp
	typename RImageBaseType::TimeStampType TimeStamp;
	m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize));
	m_InputStream.read((char*)&TimeStamp,TimeStampSize);
	Result->SetTimeStamp(TimeStamp);

	// Read the RGB values
	/*m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize));
	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 + RGBImageSize));
	m_InputStream.seekg(DataSectionStreamPos + (std::ifstream::pos_type)(FrameID*FrameSize + TimeStampSize));
	m_InputStream.read((char*)Result->GetRangeImage()->GetBufferPointer(),RangeImageSize);
	Result->GetRangeImage()->Modified();

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

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

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

	return Result;
}


//----------------------------------------------------------------------------
template<class TRImageBase>
void
RImageIO<TRImageBase>::
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 ( Superclass::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;

	// 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>
void
RImageIO<TRImageBase>::
InitTimeStamps()
{
	// Safe the stream position
	std::ifstream::pos_type DataSectionStreamPos = m_InputStream.tellg();

	// Some offsets
	long NumPixels = m_CurrentHeader->CurrCameraParameters[CameraParameters::SX]* m_CurrentHeader->CurrCameraParameters[CameraParameters::SY];
	std::ifstream::pos_type TimeStampSize = sizeof(typename RImageBaseType::TimeStampType);
	//std::ifstream::pos_type RGBImageSize = NumPixels*sizeof(RImageBaseType::RGBType);
	std::ifstream::pos_type RangeImageSize = NumPixels*sizeof(typename RImageBaseType::RangeType);
	//std::ifstream::pos_type FrameSize = TimeStampSize + RGBImageSize + RangeImageSize;
	std::ifstream::pos_type FrameSize = TimeStampSize + RangeImageSize;

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

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

} // end namespace ritk
