///////////////////////////////////////////////////////////////////////////////
// Name:			TrackIO.cpp
// Purpose:			I/O interface for track file (.trk) 
// Author:			Ruopeng Wang
// Modified by:
// Last modified:	2005.03.15
// Copyright:		(c) 2004-2005 Ruopeng Wang <rpwang@nmr.mgh.harvard.edu>
// Licence:		
//           
// Usage:
//			Sample lines to read track file -
//
//          include "TrackIO.h"
//          
//			...
//
//			CTrackReader reader;
//          TRACK_HEADER header;
//			if (!reader.Open("foo.trk", &header))
//			{
//				printf(reader.GetLastErrorMessage());
//				return;
//			}
//			...
//             
//			int cnt;
//			if (ignore_scalars)
//			{
//				while (reader.GetNextPointCount(&cnt))
//				{
//					float* pts = new float[cnt*3];
//					reader.GetNextTrackData(cnt, pts);
//					...
//		            process_point_data(...);
//					...
//					delete[] pts;
//				}
//			}
//			else
//			{
//				while (reader.GetNextPointCount(&cnt))
//				{
//					float* pts = new float[cnt*3];
//					float* scalars = new float[cnt*header.n_scalars];
//					reader.GetNextTrackData(cnt, pts, scalars);
//					...
//		            process_point_and_scalar_data(...);
//					...
//					delete[] pts;
//					delete[] scalars;
//				}
//			}
//
//			reader.Close();
//
///////////////////////////////////////////////////////////////////////////////

#include "TrackIO.h"

///// CTrackIO reference //////////////////
const char* error_message[] = 
{
	"No error", 
	"Can not open file",
	"Can not close file",
	"Can not read from file",
	"Can not write to file",
	"Not a compatible track file",
	"I/O not initialized. Call 'Open' or 'Initialize' first"
};

const char* CTrackIO::GetLastErrorMessage()
{
	return error_message[m_nErrorCode];
}

BOOL CTrackIO::GetHeader(TRACK_HEADER* header)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		return FALSE;
	}
	*header = m_header;

	return TRUE;
}

BOOL CTrackIO::Close()
{
	BOOL ret = TRUE;
	if (m_pFile)
	{
		if (fclose(m_pFile) == EOF)
		{
			m_nErrorCode = TE_CAN_NOT_CLOSE;
			ret = FALSE;
		}
		m_pFile = NULL;
	}
	if (ret)
	        m_nErrorCode = TE_NO_ERROR;

	return ret;
}

/////////////////////////////////////////

///// CTrackReader reference //////////////////

CTrackReader::CTrackReader()
{
	m_bByteSwap = FALSE;
}


BOOL CTrackReader::Open(const char* filename, TRACK_HEADER* header)
{
	Close();
	m_pFile = fopen(filename, "rb");
	if (!m_pFile)
	{
		m_nErrorCode = TE_CAN_NOT_OPEN;
		return FALSE;
	}

	fseek(m_pFile, 0, SEEK_END);
	m_nSize = ftell(m_pFile);
	fseek(m_pFile, 0, SEEK_SET);

	if (fread(&m_header, sizeof(TRACK_HEADER), 1, m_pFile) != 1)
		m_nErrorCode = TE_NOT_TRACK_FILE;
	m_header.id_string[5] = '\0';

	// could be old preliminary format from old MGH data.
	// in most cases ignored
	BOOL bOldFormat = FALSE;
	if (!m_nErrorCode && strcmp(m_header.id_string, "TRACK") != 0)
	{
		bOldFormat = TRUE;
		m_header.Initialize();
		fseek(m_pFile, 0, SEEK_SET);
		int dim[3];
		fread(dim, sizeof(int)*3, 1, m_pFile);
#ifdef DARWIN		// for big endian system
		SWAP_INT(dim, 3);
#endif
		for (int i = 0; i < 3; i++)
			m_header.dim[i] = dim[i];
#ifdef DARWIN
		SWAP_SHORT(m_header.dim, 3);
#endif
		fread(m_header.voxel_size, sizeof(float)*3, 1, m_pFile);
		m_header.n_scalars = 0;
#ifdef DARWIN
		m_bByteSwap = TRUE;
#endif
	}
	else
	{
	        if (m_header.hdr_size == 0)  // should not happen, for test only
		{
#ifdef DARWIN
			m_bByteSwap = TRUE;
#else 
			m_bByteSwap = FALSE;
#endif
		}
		else
			m_bByteSwap = (m_header.hdr_size != sizeof(struct TRACK_HEADER));
	}
	
	if (m_bByteSwap)
	{
		m_header.ByteSwap();
		m_header.hdr_size = sizeof(struct TRACK_HEADER);
	}

	if (bOldFormat)
	{
		if (fmin(m_header.dim, 3) <= 0 || fmin(m_header.voxel_size, 3) <= 0.00001 ||
			fmax(m_header.dim, 3) > 10000 || fmax(m_header.voxel_size, 3) > 100)
			m_nErrorCode = TE_NOT_TRACK_FILE;
	}

	if (header)
		*header = m_header;

	m_bOldFormat = bOldFormat;

	if (m_nErrorCode != TE_NO_ERROR)
	{
		Close();
		return FALSE;
	}

	return TRUE;
}

// Get progess in percentage
int CTrackReader::GetProgress()
{
	if (m_pFile && m_nSize > 0)
		return (int)((double)(ftell(m_pFile)) * 100 / (m_nSize));
	return 0;
}

BOOL CTrackReader::GetNextPointCount(int* n)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		n = 0;
		return FALSE;
	}
	if (fread(n, sizeof(int), 1, m_pFile) != 1)
		m_nErrorCode = TE_CAN_NOT_READ;
	else
		m_nErrorCode = TE_NO_ERROR;

	if (m_bByteSwap)
		SWAP_INT(*n);
	
	return m_nErrorCode == TE_NO_ERROR;
}


// Attention: data buffer has to been pre-allocated. The following two routines 
// do not allocate memory. 
// GetNextRawData(...) read in together the point coordinates and scalars
// GetNextTrackData(...) read in point coords and scalars in seperated buffers
// if scalars buffer is not given, only point coords are read.
// These two routines can only be called once each time after calling GetNextPointCount().
BOOL CTrackReader::GetNextRawData(int nCount, float* data)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		return FALSE;
	}
	int nSize = (3+m_header.n_scalars)*nCount;
	if (fread(data, sizeof(float)*nSize, 1, m_pFile) != 1)
		m_nErrorCode = TE_CAN_NOT_READ;
	else
		m_nErrorCode = TE_NO_ERROR;

	if (!m_nErrorCode && m_bByteSwap)
		SWAP_FLOAT(data, nSize);

	return m_nErrorCode == TE_NO_ERROR;
}

BOOL CTrackReader::GetNextTrackData(int nCount, float* pt_data, float* scalars)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		return FALSE;
	}
	m_nErrorCode = TE_NO_ERROR;
	
	if (m_header.n_scalars == 0)
		return GetNextRawData(nCount, pt_data);

	for (int i = 0; i < nCount; i++)
	{
		if (fread(pt_data+i*3, sizeof(float)*3, 1, m_pFile) != 1)
		{
			m_nErrorCode = TE_CAN_NOT_READ;
			break;
		}
		if (scalars)
		{
			if (fread(scalars+i*m_header.n_scalars, sizeof(float)*m_header.n_scalars, 1, m_pFile) != 1)
			{
				m_nErrorCode = TE_CAN_NOT_READ;
				break;
			}
		}
		else
		{
			fseek(m_pFile, sizeof(float)*m_header.n_scalars, SEEK_CUR);
		}
	}
	if (!m_nErrorCode && m_bByteSwap)
	{
		SWAP_FLOAT(pt_data, nCount*3);
		if (scalars)
			SWAP_FLOAT(scalars, nCount*m_header.n_scalars);
	}

	return m_nErrorCode == TE_NO_ERROR;
}

// Static function. Get header info from a given track file directly
BOOL CTrackReader::GetHeaderOnly(const char* filename, TRACK_HEADER *header)
{
	CTrackReader reader;
	BOOL ret = reader.Open(filename, header);
	reader.Close();

	return ret;
}

// if number of tracks was not recorded in the header, this routine will 
// take longer to excute as it will go through the whole file
BOOL CTrackReader::GetNumberOfTracks(int* cnt)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		*cnt = 0;
		return FALSE;
	}
	if (m_header.n_count == 0)
	{
		long pos = ftell(m_pFile);
		fseek(m_pFile, sizeof(TRACK_HEADER), SEEK_SET);
		
		int n;
		*cnt = 0;
		while (GetNextPointCount(&n))
		{
			*cnt += 1;
			fseek(m_pFile, sizeof(float)*n*(3+m_header.n_scalars), SEEK_CUR);
		}
		fseek(m_pFile, pos, SEEK_SET);
		m_header.n_count = *cnt;
	}
	else
	{
		*cnt = m_header.n_count;
	}
	m_nErrorCode = TE_NO_ERROR;

	return TRUE;
}

int CTrackReader::GetNumberOfTracks()
{
	int n;
	GetNumberOfTracks(&n);
	return n;
}

///////////////////////////////////////////////

///// CTrackReader reference //////////////////

// One of the Initializers must be called before WriteNextTrackData()
BOOL CTrackWriter::Initialize(const char* filename, short int* dim, float* voxel_size, float* origin,
		short int n_scalars)
{
	TRACK_HEADER header(dim, voxel_size, origin, n_scalars);

	return Initialize(filename, header);
}

BOOL CTrackWriter::Initialize(const char* filename, int* dim, float* voxel_size, float* origin,
		short int n_scalars)
{
	short int ndim[3];
	for (int i = 0; i < 3; i++)
		ndim[i] = dim[i];

	return Initialize(filename, ndim, voxel_size, origin, n_scalars);
}

BOOL CTrackWriter::Initialize(const char* filename, TRACK_HEADER header)
{
	Close();
	m_pFile = fopen(filename, "wb");
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		return FALSE;
	}

	if (header.hdr_size != sizeof(TRACK_HEADER))
		header.ByteSwap();
	m_header = header;
	m_header.hdr_size = sizeof(struct TRACK_HEADER);

	if (fwrite(&m_header, sizeof(struct TRACK_HEADER), 1, m_pFile) == 1)
		m_nErrorCode = TE_NO_ERROR;
	else
		m_nErrorCode = TE_CAN_NOT_WRITE;

	m_header.n_count = 0;
	return m_nErrorCode == TE_NO_ERROR;
}


// data must include scalars if n_scalars is not 0
BOOL CTrackWriter::WriteNextTrack(int ncount, float* data)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		return FALSE;
	}
	int n = ncount;
	m_nErrorCode = TE_NO_ERROR;
	
	if (fwrite(&n, sizeof(int), 1, m_pFile) != 1)
		m_nErrorCode = TE_CAN_NOT_WRITE;
	long nSize = ncount*(3+m_header.n_scalars)*sizeof(float);

	if (fwrite(data, nSize, 1, m_pFile) != 1)
		m_nErrorCode = TE_CAN_NOT_WRITE;

	if (!m_nErrorCode)
		m_header.n_count ++;
	return m_nErrorCode == TE_NO_ERROR;
}

BOOL CTrackWriter::UpdateHeader(TRACK_HEADER header)
{
	if (!m_pFile)
	{
		m_nErrorCode = TE_NOT_INITIALIZED;
		return FALSE;
	}

	if (header.hdr_size != sizeof(TRACK_HEADER))
		header.ByteSwap();
	
	m_header = header;
	m_header.hdr_size = sizeof(struct TRACK_HEADER);

	long pos = ftell(m_pFile);
	fseek(m_pFile, 0, SEEK_SET);
	
	if (fwrite(&m_header, sizeof(struct TRACK_HEADER), 1, m_pFile) != 1)
		m_nErrorCode = TE_CAN_NOT_WRITE;
	else
		m_nErrorCode = TE_NO_ERROR;
	fseek(m_pFile, pos, SEEK_SET);

	return m_nErrorCode == TE_NO_ERROR;
}

BOOL CTrackWriter::Close()
{
	return UpdateHeader(m_header) && CTrackIO::Close();
}

// Static function. Write the given header to a existing track file. 
BOOL CTrackWriter::UpdateHeader(const char* filename, TRACK_HEADER header)
{
	TRACK_HEADER hdr;
	FILE* fp = fopen(filename, "r+b");
	if (!fp)
		return FALSE;

	if (fread(&hdr, sizeof(TRACK_HEADER), 1, fp) != 1)
	{
		fclose(fp);
		return FALSE;
	}

	BOOL bswap = (hdr.hdr_size != sizeof(struct TRACK_HEADER));
	hdr = header;
	hdr.hdr_size = sizeof(struct TRACK_HEADER);

	if (bswap)
		hdr.ByteSwap();

	fseek(fp, 0, SEEK_SET);
	if (fwrite(&hdr, sizeof(TRACK_HEADER), 1, fp) != 1)
	{
		fclose(fp);
		return FALSE;
	}
	fclose(fp);

	return TRUE;
}