#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkRAWImageIO.h"
#include "itkImageRegionIteratorWithIndex.h"

typedef itk::Image<unsigned char, 3> VOLType;

VOLType::Pointer rts_itkNewVOL(int x, int y, int z)
{
	//Create a new image and allocate the desired amount of space
	//typedef itk::Image<T, 3> ImageType;
	VOLType::Pointer result = VOLType::New();
	VOLType::RegionType region;
	VOLType::SizeType size = {x, y, z};
	region.SetSize(size);
	VOLType::IndexType index = {0, 0, 0};
	region.SetIndex(index);
	result->SetRegions(region);
	result->Allocate();

	return result;

}

VOLType::Pointer rts_itkLoadVOL(const char* filename, unsigned int max_slices = 512)
{
	//first load the header information
	ifstream infile(filename, ios::in | ios::binary);	//create the files stream
	if(!infile)
		return VOLType::New();

	unsigned int size_x, size_y, size_z;				//create variables to store the size of the data set
	//load the dimensions of the data set
	infile.read((char*)&size_x, sizeof(int));			//load the file header
	infile.read((char*)&size_y, sizeof(int));
	infile.read((char*)&size_z, sizeof(int));

	//close the file
	infile.close();


	VOLType::Pointer volume;


	//create the reader
	typedef itk::ImageFileReader<VOLType> ReaderType;
	typedef itk::RawImageIO<unsigned char, 3> RawType;

	ReaderType::Pointer reader = ReaderType::New();
	RawType::Pointer rawIO = RawType::New();

	

	reader->SetImageIO(rawIO);
	reader->SetFileName(filename);

	//set the raw IO parameters
	rawIO->SetFileTypeToBinary();
	

	//set the format of the RAW file
	rawIO->SetHeaderSize(12);
	rawIO->SetDimensions(0, size_x);
	rawIO->SetDimensions(1, size_y);

	if(size_z < max_slices)
		rawIO->SetDimensions(2, size_z);
	else
		rawIO->SetDimensions(2, max_slices);

	//read the file
	try
	{
		reader->Update();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}

	
	volume = reader->GetOutput();

	return volume;
}
VOLType::Pointer rts_itkLoadRAW(const char* filename, unsigned int size_x, unsigned int size_y, unsigned int size_z)
{
	//first load the header information
	ifstream infile(filename, ios::in | ios::binary);	//create the files stream
	if(!infile)
		return VOLType::New();

	

	//close the file
	infile.close();


	VOLType::Pointer volume;


	//create the reader
	typedef itk::ImageFileReader<VOLType> ReaderType;
	typedef itk::RawImageIO<unsigned char, 3> RawType;

	ReaderType::Pointer reader = ReaderType::New();
	RawType::Pointer rawIO = RawType::New();

	

	reader->SetImageIO(rawIO);
	reader->SetFileName(filename);

	//set the raw IO parameters
	rawIO->SetFileTypeToBinary();
	

	//set the format of the RAW file
	rawIO->SetHeaderSize(0);
	rawIO->SetDimensions(0, size_x);
	rawIO->SetDimensions(1, size_y);
	rawIO->SetDimensions(2, size_z);


	//read the file
	try
	{
		reader->Update();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}

	
	volume = reader->GetOutput();

	return volume;
}
void rts_itkSaveVOL(VOLType::Pointer itk_image, const char* filename)
{
	typedef itk::ImageRegionIteratorWithIndex<VOLType> IteratorType;
	IteratorType i(itk_image, itk_image->GetLargestPossibleRegion());

	//create variables for size and position
	VOLType::IndexType index;
	VOLType::SizeType size;
	//get the volume size
	size = itk_image->GetLargestPossibleRegion().GetSize();
	unsigned char* buffer = new unsigned char[size[0]*size[1]*size[2]];

	for(i.GoToBegin(); !i.IsAtEnd(); ++i)
	{
		index = i.GetIndex();
		buffer[index[0] + size[0] * (index[1] + index[2] * size[1])] = i.Get();
	}


	ofstream outfile(filename, ios::out | ios::binary);	//create the binary file stream

	//write the volume size to the file
	outfile.write((char*)&size[0], sizeof(int));
	outfile.write((char*)&size[1], sizeof(int));
	outfile.write((char*)&size[2], sizeof(int));

	outfile.write((char*)buffer, sizeof(char)*size[0]*size[1]*size[2]);
	outfile.close();
}