#ifndef RTS_ITKVOLUME_H
#define RTS_ITKVOLUME_H

#include <WTypes.h>
#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkRawImageIO.h"
#include "itkImageRegionIterator.h"
#include <string>
#include "rtsFilename.h"
#include "rtsDTGrid3D.h"
using namespace std;

#include "itkRescaleIntensityImageFilter.h"

#include "rts_itkImage.h"

void rtsProgressBar(int percent);

template <class T> class rts_itkVolume
{
public:
	typedef int uint;
	typedef itk::Image<T, 3> ITKVolumeType;	
	typedef typename ITKVolumeType::Pointer PointerType;
	typedef itk::Image<T, 2> ITKSliceType;
	
	//typedef itk::Image<typename T, 3> ITKVolumeType;
private:
	//pointer to the data
	PointerType imageData;

	char* UnicodeToANSI(wchar_t* unicodestr)
	{
		int a = wcslen(unicodestr)+1;
		char *ansistr = new char[a];
		::WideCharToMultiByte(CP_ACP, 
								0, 
								unicodestr, 
								-1, 
								ansistr, 
								a, 
								NULL, 
								NULL);
		return ansistr;

	}

	BSTR ANSIToUnicode(const char* ansistr)
	{
		/*On a side note:  This BSTR crap isn't the way to go.
		Try to use wchar_t (Unicode) where possible.  Already changed it in
		the above function because it was causing a bug.*/

		int a = strlen(ansistr);
		BSTR unicodestr = SysAllocStringLen(NULL, a);
		::MultiByteToWideChar(CP_ACP, 0, ansistr, a, unicodestr, a);
		return unicodestr;
	}

public:
	//construct an implicit function with a size of 1
	rts_itkVolume();			///<Create an empty implicit function
	//construct an implicit function of the specified resolution
	rts_itkVolume(uint res_x, uint res_y, uint res_z);	///<Create an implicit function with the specified resolution
	void Init(uint res_x, uint res_y, uint res_z);
	void SetSpacing(float voxel_x, float voxel_y, float voxel_z);
	void Clear();

	
	//loading/saving data to disk
	void LoadRAW(uint header_size, uint data_x, uint data_y, uint data_z, string filename);	///<Loads RAW data from a file with the specified header size and data size.
	void SaveRAW(string filename);							///<Save the data as RAW data to disk.
	void LoadVOL(string filename);							///<Load a VOL file from disk.
	void SaveVOL(string filename);							///<Save a VOL file to disk.
	void LoadRAV(string filename);							///<Load a VOL file from disk.
	void SaveRAV(string filename);							///<Save a VOL file to disk.
	void LoadStack(string directory, string filename_mask, int max_files = -1);
	void LoadSlice(string filename, int slice);

	void InsertSlice(uint z_position, typename ITKSliceType::Pointer slice);

	
	//data access methods
	T operator ()(uint x, uint y, uint z);
	PointerType GetITKPointer(){return imageData;}
	T* GetBits(){return imageData->GetBufferPointer();}
	T GetMin();
	T GetMax();

	void SetITKPointer(PointerType pointer){imageData = pointer;}
	uint DimX(){return imageData->GetLargestPossibleRegion().GetSize()[0];}
	uint DimY(){return imageData->GetLargestPossibleRegion().GetSize()[1];}
	uint DimZ(){return imageData->GetLargestPossibleRegion().GetSize()[2];}

	float VoxelSizeX(){return imageData->GetSpacing()[0];}
	float VoxelSizeY(){return imageData->GetSpacing()[1];}
	float VoxelSizeZ(){return imageData->GetSpacing()[2];}
	void SetPixel(uint x, uint y, uint z, T value);
	void SetSlice(rts_itkImage<T> image, int slice);
	T GetPixel(uint x, uint y, uint z);
	rts_itkImage<T> GetSlice(int slice);
	rts_itkVolume<T> SubVolume(uint px, uint py, uint pz, uint sx, uint sy, uint sz);

	//mathematica operations
	T operator=(T rhs);
	rts_itkVolume<T> operator-(T rhs);
	rts_itkVolume<T> operator+(T rhs);
	rts_itkVolume<T> operator-(rts_itkVolume<T> rhs);

	void Print(ostream &os);

	//simple data processing
	void Rescale(T new_min, T new_max);
	void Binary(T threshold, T true_value);		///<Turns the image into a binary image based on a threshold value T.  All values below T are set to 0, all values above are set to true_value.

	void InsertDTGrid(rtsDTGrid3D<T>* grid, bool debug = false);
};

template <class T>
rts_itkVolume<T> rts_itkVolume<T>::SubVolume(uint px, uint py, uint pz, uint sx, uint sy, uint sz)
{
	/*This function returns the sub-volume specified by a corner point and size*/

	//first make sure that the sub-volume is contained within the image
	if(px < 0) {sx += px; px = 0;}
	if(py < 0) {sy += py; py = 0;}
	if(pz < 0) {sz += pz; pz = 0;}
	if(px >= DimX()) px = DimX() - 1;
	if(py >= DimY()) py = DimY() - 1;
	if(pz >= DimZ()) pz = DimZ() - 1;
	if(px + sx > DimX()) sx = DimX() - px - 1;
	if(py + sy > DimY()) sy = DimY() - py - 1;
	if(pz + sz > DimZ()) sz = DimZ() - pz - 1;

	//cout<<px<<","<<py<<","<<pz<<","<<sx<<","<<sy<<","<<sz<<endl;

	//create the resulting volume
	rts_itkVolume destination(sx, sy, sz);

	//get an iterator for the destination volume
	itk::ImageRegionIterator<ITKVolumeType> d(destination.imageData, destination.imageData->GetLargestPossibleRegion());

	//get an iterator for the source region
	ITKVolumeType::RegionType region;
	ITKVolumeType::IndexType index;
	ITKVolumeType::SizeType size;
	index[0] = px; index[1] = py; index[2] = pz;
	size[0] = sx; size[1] = sy; size[2] = sz;
	region.SetIndex(index);
	region.SetSize(size);
	itk::ImageRegionIterator<ITKVolumeType> s(imageData, region);

	//copy the data from source region to destination volume
	for(s.GoToBegin(), d.GoToBegin();
		!s.IsAtEnd();
		s++, d++)
		d.Set(s.Get());

	return destination;

}

template <class T>
T rts_itkVolume<T>::operator=(T rhs)
{
	/*itk::ImageRegionIterator<ITKVolumeType> i(imageData, imageData->GetLargestPossibleRegion());
	
	for(i.GoToBegin(); !i.IsAtEnd(); i++)
		i.Set(rhs);
*/
	imageData->FillBuffer(rhs);
	return rhs;
}
template <class T>
rts_itkVolume<T> rts_itkVolume<T>::operator-(T rhs)
{
	rts_itkVolume<T> result(DimX(), DimY(), DimZ());

	itk::ImageRegionIterator<ITKVolumeType> w(result.imageData, result.imageData->GetLargestPossibleRegion());
	itk::ImageRegionIterator<ITKVolumeType> r(imageData, imageData->GetLargestPossibleRegion());

	for(r.GoToBegin(), w.GoToBegin(); !r.IsAtEnd(); r++, w++)
		w.Set(r.Get() - rhs);

	return result;
}

template <class T>
rts_itkVolume<T> rts_itkVolume<T>::operator+(T rhs)
{
	rts_itkVolume<T> result(DimX(), DimY(), DimZ());

	itk::ImageRegionIterator<ITKVolumeType> w(result.imageData, result.imageData->GetLargestPossibleRegion());
	itk::ImageRegionIterator<ITKVolumeType> r(imageData, imageData->GetLargestPossibleRegion());

	for(r.GoToBegin(), w.GoToBegin(); !r.IsAtEnd(); r++, w++)
		w.Set(r.Get() + rhs);

	return result;
}

template <class T>
rts_itkVolume<T> rts_itkVolume<T>::operator-(rts_itkVolume<T> rhs)
{
	rts_itkVolume<T> result(DimX(), DimY(), DimZ());

	itk::ImageRegionIterator<ITKVolumeType> w(result.imageData, result.imageData->GetLargestPossibleRegion());
	itk::ImageRegionIterator<ITKVolumeType> r(rhs.imageData, rhs.imageData->GetLargestPossibleRegion());
	itk::ImageRegionIterator<ITKVolumeType> l(imageData, imageData->GetLargestPossibleRegion());

	for(r.GoToBegin(), l.GoToBegin(), w.GoToBegin(); 
		!r.IsAtEnd(); 
		r++, l++, w++)
		w.Set(l.Get() - r.Get());

	return result;
}

template <class T>
void rts_itkVolume<T>::Print(ostream &os)
{
	ITKVolumeType::SizeType size;
	size = imageData->GetLargestPossibleRegion().GetSize();

	int x, y, z;
	for(z=0; z<size[2]; z++)
	{
		os<<"z = "<<z<<":"<<endl;
		for(y=0; y<size[1]; y++)
		{
			for(x=0; x<size[0]; x++)
				os<<GetPixel(x, y, z)<<" ";
			os<<endl;
		}
	}
}

template <class T>
rts_itkVolume<T>::rts_itkVolume()
{
	imageData = ITKVolumeType::New();
}

template <class T>
rts_itkVolume<T>::rts_itkVolume(uint res_x, uint res_y, uint res_z)
{
	imageData = ITKVolumeType::New();

	itk::Image<T, 3>::RegionType region;
	itk::Image<T, 3>::SizeType size;
	size[0] = res_x;
	size[1] = res_y;
	size[2] = res_z;

	region.SetSize(size);
	imageData->SetRegions(region);
	
	try
	{
		imageData->Allocate();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}
		
}

template <class T>
void rts_itkVolume<T>::Init(uint res_x, uint res_y, uint res_z)
{
	itk::Image<T, 3>::RegionType region;
	itk::Image<T, 3>::SizeType size;
	size[0] = res_x;
	size[1] = res_y;
	size[2] = res_z;

	region.SetSize(size);
	imageData->SetRegions(region);

	
	try
	{
		imageData->Allocate();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}
}

template <class T>
void rts_itkVolume<T>::SetSpacing(float voxel_x, float voxel_y, float voxel_z)
{
	float spacing[3];
	spacing[0] = voxel_x; spacing[1]= voxel_y; spacing[2] = voxel_z;
	imageData->SetSpacing(spacing);
}


template <class T>
inline T rts_itkVolume<T>::GetPixel(uint x, uint y, uint z)
{
	T result;
	if(x < 0 || x >= DimX() || y < 0 || y >= DimY() || z < 0 || z >= DimZ())
	{
		memset(&result, 0, sizeof(T));
		return result;
	}

	itk::Image<T, 3>::IndexType index;
	index[0] = x; index[1] = y; index[2] = z;
	return imageData->GetPixel(index);
}

template <class T>
inline T rts_itkVolume<T>::GetMin()
{
	int ix, iy, iz;
	T currentMin = GetPixel(0, 0, 0);
	T n;
	for(ix=0; ix<DimX(); ix++)
		for(iy=0; iy<DimY(); iy++)
			for(iz=0; iz<DimZ(); iz++)
			{
				n = GetPixel(ix, iy, iz);
				if(n < currentMin)
					currentMin = n;
			}
	return currentMin;
}

template <class T>
inline T rts_itkVolume<T>::GetMax()
{
	int ix, iy, iz;
	T currentMax = GetPixel(0, 0, 0);
	T n;
	for(ix=0; ix<DimX(); ix++)
		for(iy=0; iy<DimY(); iy++)
			for(iz=0; iz<DimZ(); iz++)
			{
				n = GetPixel(ix, iy, iz);
				if(n > currentMax)
					currentMax = n;
			}
	return currentMax;
}

template <class T>
inline T rts_itkVolume<T>::operator ()(uint x, uint y, uint z)
{
	return GetPixel(x, y, z);
}



template <class T>
void rts_itkVolume<T>::LoadRAW(uint header_size, uint size_x,
							   uint size_y, uint size_z, string filename)
{
	ITKVolumeType::RegionType region;
	ITKVolumeType::SizeType size;
	size[0] = size_x;
	size[1] = size_y;
	size[2] = size_z;
	region.SetSize(size);
	imageData->SetRegions(region);

	try
	{
		imageData->Allocate();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}

	
	//cout<<imageData<<endl;

	ifstream infile(filename.c_str(), ios::out | ios::binary);	//create the files stream
	if(!infile)
		return;
	//first read the header
	char* header = new char[header_size];
	infile.read((char*)header, header_size);

	//now read the actual data
	infile.read((char*)imageData->GetBufferPointer(), DimX()*DimY()*DimZ()*sizeof(T));
	
	infile.close();
}

template <class T>
void rts_itkVolume<T>::LoadRAV(string filename)
{
	
	//open the file
	ifstream infile(filename.c_str(), ios::out | ios::binary);	//create the files stream
	if(!infile)
		return;

	//first read the header
	int dimx, dimy, dimz, bytesize;
	infile.read((char*)&dimx, sizeof(int));
	infile.read((char*)&dimy, sizeof(int));
	infile.read((char*)&dimz, sizeof(int));
	infile.read((char*)&bytesize, sizeof(int));

	//allocate space for the image
	ITKVolumeType::RegionType region;
	ITKVolumeType::SizeType size;
	size[0] = dimx;
	size[1] = dimy;
	size[2] = dimz;
	region.SetSize(size);
	imageData->SetRegions(region);
	try
	{
		imageData->Allocate();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}
	

	//now read the actual data
	infile.read((char*)imageData->GetBufferPointer(), DimX()*DimY()*DimZ()*sizeof(T));
	
	infile.close();	
}

template <class T>
void rts_itkVolume<T>::LoadVOL(string filename)
{
	
	//open the file
	ifstream infile(filename.c_str(), ios::out | ios::binary);	//create the files stream
	if(!infile)
		return;

	//first read the header
	int dimx, dimy, dimz, bytesize;
	infile.read((char*)&dimx, sizeof(int));
	infile.read((char*)&dimy, sizeof(int));
	infile.read((char*)&dimz, sizeof(int));

	//allocate space for the image
	ITKVolumeType::RegionType region;
	ITKVolumeType::SizeType size;
	size[0] = dimx;
	size[1] = dimy;
	size[2] = dimz;
	region.SetSize(size);
	imageData->SetRegions(region);
	try
	{
		imageData->Allocate();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}
	

	//now read the actual data
	infile.read((char*)imageData->GetBufferPointer(), DimX()*DimY()*DimZ()*sizeof(T));
	
	infile.close();	
}

template <class T>
void rts_itkVolume<T>::LoadStack(string directory, string filename_mask, int max_files = -1)
{
	/*Load a list of file names based on the user-specified mask*/
	//create the complete mask
	if(directory[directory.size() - 1] != '/')
		directory += '/';

	string full_mask = directory;	
	full_mask += filename_mask;
	
	//create a temporary vector to store the file names
	vector<rtsFilename> fileList;
	//create a handle for the search
	HANDLE search_handle;
	WIN32_FIND_DATA found_data;

	//start the search
	BSTR unicodestr = ANSIToUnicode(full_mask.c_str());
	search_handle = FindFirstFile(unicodestr, &found_data);
	::SysFreeString(unicodestr);

	//if there are no files, return
	if(search_handle == INVALID_HANDLE_VALUE)
		return;

	string new_file;
	char* ansistr = UnicodeToANSI(found_data.cFileName);
	//char* ansistr = (char*)found_data.cFileName;
	//save the first filename
	new_file = directory;
	new_file += ansistr;
	fileList.push_back(new_file);
	delete[] ansistr;
		
	//now loop through the rest of the files
	while(FindNextFile(search_handle, &found_data))
	{
		char* ansistr = UnicodeToANSI(found_data.cFileName);
		new_file = directory;
		new_file += ansistr;
		//cout<<new_file<<endl;
		fileList.push_back(new_file);
		delete[] ansistr;
	}



	/*Load the first file in order to get volume information*/
	typedef itk::ImageFileReader<ITKSliceType> ReaderType;
	ReaderType::Pointer reader = ReaderType::New();
	reader->SetFileName(fileList[0].getString());
	ITKSliceType::Pointer image = reader->GetOutput();
	
	try
	{
		reader->Update();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}

	/*Insert each image into the volume*/
	int i;
	cout<<"Files found: "<<fileList.size()<<endl;
	if(max_files == -1)
		max_files = fileList.size();
	else
		max_files = min<int>(fileList.size(), max_files);
	cout<<"Loading: "<<max_files<<endl;
	//set the max number of files based on user input and number of files found

	/*Set up the volume based on information from the first file*/
	ITKSliceType::SizeType size = image->GetLargestPossibleRegion().GetSize();	//get the slice size
	ITKVolumeType::SizeType volSize;	//compute the volume size
	volSize[0] = size[0];
	volSize[1] = size[1];
	volSize[2] = max_files;

	imageData = ITKVolumeType::New();	//create the volume
	ITKVolumeType::RegionType region;
	region.SetSize(volSize);
	imageData->SetRegions(region);
	imageData->Allocate();				//allocate space



	for(i=0; i<max_files; i++)
	{
		reader = ReaderType::New();
		reader->SetFileName(fileList[i].getString());
		ITKSliceType::Pointer slice = reader->GetOutput();
		reader->Update();
		InsertSlice(i, slice);
		//cout<<i+1<<"/"<<fileList.size()<<endl;
		rtsProgressBar((float)(i+1)/(float)max_files*100);
	}	
}

template <class T>
void rts_itkVolume<T>::LoadSlice(string filename, int slice)
{
	//load the specified image
	typedef itk::ImageFileReader<ITKSliceType> ReaderType;
	ReaderType::Pointer reader = ReaderType::New();
	reader->SetFileName(filename.c_str());
	ITKSliceType::Pointer image = reader->GetOutput();
	
	try
	{
		reader->Update();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}

	//insert the slice into the volume
	InsertSlice(slice, image);

}

template <class T>
rts_itkImage<T> rts_itkVolume<T>::GetSlice(int slice)
{
	//create a new image
	rts_itkImage<T> outSlice(DimX(), DimY());

	int x, y;
	for(x=0; x<DimX(); x++)
		for(y=0; y<DimY(); y++)
		{
			outSlice.SetPixel(x, y, GetPixel(x, y, slice));
		}

	return outSlice;
}

template <class T>
void rts_itkVolume<T>::SetSlice(rts_itkImage<T> image, int slice)
{
	int x, y;
	for(x=0; x<DimX(); x++)
		for(y=0; y<DimY(); y++)
		{
			SetPixel(x, y, slice, image.GetPixel(x, y));
		}
}

template <class T>
void rts_itkVolume<T>::SaveRAV(string filename)
{
	ofstream outfile(filename.c_str(), ios::out | ios::binary);	//create the files stream
	if(!outfile)
		return;

	int dimx = DimX();
	int dimy = DimY();
	int dimz = DimZ();
	int bytesize = sizeof(T);

	outfile.write((char*)&dimx, sizeof(int));
	outfile.write((char*)&dimy, sizeof(int));
	outfile.write((char*)&dimz, sizeof(int));
	outfile.write((char*)&bytesize, sizeof(int));

	outfile.write((char*)imageData->GetBufferPointer(), DimX()*DimY()*DimZ()*sizeof(T));
	outfile.close();
	
}

template <class T>
void rts_itkVolume<T>::SaveVOL(string filename)
{
	ofstream outfile(filename.c_str(), ios::out | ios::binary);	//create the files stream
	if(!outfile)
		return;

	int dimx = DimX();
	int dimy = DimY();
	int dimz = DimZ();

	outfile.write((char*)&dimx, sizeof(int));
	outfile.write((char*)&dimy, sizeof(int));
	outfile.write((char*)&dimz, sizeof(int));

	outfile.write((char*)imageData->GetBufferPointer(), DimX()*DimY()*DimZ()*sizeof(T));
	outfile.close();
	
}

template <class T>
void rts_itkVolume<T>::InsertSlice(uint z_position, typename ITKSliceType::Pointer slice)
{
	itk::Image<T>::SizeType size = slice->GetLargestPossibleRegion().GetSize();
	itk::Image<T>::IndexType index;
	int x, y;
	for(x=0; x<size[0]; x++)
		for(y=0; y<size[1]; y++)
		{
			index[0] = x; index[1] = y;
			SetPixel(x, y, z_position, slice->GetPixel(index));
		}



}

template <class T>
void rts_itkVolume<T>::SaveRAW(string filename)
{
	ofstream outfile(filename.c_str(), ios::out | ios::binary);	//create the files stream
	if(!outfile)
		return;

	outfile.write((char*)imageData->GetBufferPointer(), DimX()*DimY()*DimZ()*sizeof(T));
	outfile.close();
}

template <class T>
void rts_itkVolume<T>::SetPixel(uint x, uint y, uint z, T value)
{
	ITKVolumeType::IndexType index;
	index[0] = x;
	index[1] = y;
	index[2] = z;
	imageData->SetPixel(index, value);
}

template <class T>
void rts_itkVolume<T>::Rescale(T new_min, T new_max)
{
	typedef itk::RescaleIntensityImageFilter<ITKVolumeType, ITKVolumeType> RescaleFilterType;
	RescaleFilterType::Pointer rescaleFilter = RescaleFilterType::New();
	rescaleFilter->SetOutputMinimum(new_min);
	rescaleFilter->SetOutputMaximum(new_max);
	rescaleFilter->SetInput(imageData);
	rescaleFilter->Update();

	imageData = rescaleFilter->GetOutput();
}

template <class T>
void rts_itkVolume<T>::Binary(T threshold, T true_value)
{
	/**
	This function converts an implicit function into a binary or characteristic function describing the solid represented by the level
	set at isovalue "threshold".  All values below threshold are set to zero while all values above threshold are set to the specified
	"true_value".  In order to use this function, the data type T must be able to be set to 0.
	**/
	int max_index = DimX() * DimY() * DimZ();	//find the size of the data array
	int i;
	T* ptrBits = imageData->GetBufferPointer();
	for(i=0; i<max_index; i++)
		if(ptrBits[i] >= threshold)
			ptrBits[i] = true_value;
		else
			ptrBits[i] = 0;
}


template<typename T>
void rts_itkVolume<T>::InsertDTGrid(rtsDTGrid3D<T>* grid, bool debug = false)
{
	//find the extents of the DT Grid
	int min_x, min_y, min_z, max_x, max_y, max_z;
	grid->getBounds(min_x, min_y, min_z, max_x, max_y, max_z);

	//allocate the appropriate amount of space
	ITKVolumeType::RegionType region;
	ITKVolumeType::SizeType size;
	size[0] = (max_x - min_x) + 1;
	size[1] = (max_y - min_y) + 1;
	size[2] = (max_z - min_z) + 1;
	region.SetSize(size);
	imageData->SetRegions(region);

	try
	{
		imageData->Allocate();
	}
	catch(itk::ExceptionObject & exp)
	{
		std::cout<<exp<<std::endl;
	}

	//fill the image with the background color
	imageData->FillBuffer(grid->background);

	//copy values from the DT Grid to the image
	ITKVolumeType::IndexType index;		//create an image index
	rtsDTGrid3D<T>::iterator i;			//create an iterator

	for(i = grid->begin(); i!= grid->after(); i++)
	{
		index[0] = i.X1() - min_x;
		index[1] = i.X2() - min_y;
		index[2] = i.X3() - min_z;
		imageData->SetPixel(index, i.Value());

		//cout<<index[0]<<","<<index[1]<<": "<<i.Value()<<endl;
		
	}
}
#endif