#ifndef _RTS_ITK_IMAGE_H
#define _RTS_ITK_IMAGE_H

#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkCastImageFilter.h"
#include "itkRescaleIntensityImageFilter.h"
#include "rtsDTGrid2D.h"

#include <string>

using namespace std;



template<typename T>
class rts_itkImage
{
public:
	typedef itk::Image<T> ITKImageType;
	typedef typename ITKImageType::Pointer ITKPointerType;
	ITKPointerType image_source;

public:
	rts_itkImage();
	rts_itkImage(int res_x, int res_y);
	void Allocate(int res_x, int res_y);
	void LoadImage(string filename);
	void LoadRAW(int header_size, int res_x, int res_y, string filename);
	void SaveImage(string filename);
	void SaveRAW(string filename);
	void CastAndSaveImage(string filename);
	void RescaleImage(T min, T max);
	void InsertDTGrid(rtsDTGrid2D<T>* grid, bool debug = 0);
	T* GetPointer();
	rts_itkImage<T> SubImage(int px, int py, int sx, int sy);

	void SetPixel(int x, int y, T value);
	T GetPixel(int x, int y);

	//element-wise operators
	rts_itkImage<T> operator+(T rhs);

	int DimX(){return image_source->GetLargestPossibleRegion().GetSize()[0];}
	int DimY(){return image_source->GetLargestPossibleRegion().GetSize()[1];}


};

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

	int x, y;
	for(x=0; x<DimX(); x++)
		for(y=0; y<DimY(); y++)
			result.SetPixel(x, y, GetPixel(x, y) + rhs);

	return result;
}

template <typename T>
void rts_itkImage<T>::SetPixel(int x, int y, T value)
{
	ITKImageType::IndexType index;
	index[0] = x;
	index[1] = y;
	image_source->SetPixel(index, value);
}

template <typename T>
T rts_itkImage<T>::GetPixel(int x, int y)
{
	ITKImageType::IndexType index;
	index[0] = x;
	index[1] = y;
	return image_source->GetPixel(index);
}

template <class T>
rts_itkImage<T>::rts_itkImage()
{
	if(image_source.IsNotNull())
		image_source = NULL;
	image_source = ITKImageType::New();
}

template <class T>
rts_itkImage<T>::rts_itkImage(int res_x, int res_y)
{
	Allocate(res_x, res_y);		
}

template<typename T>
void rts_itkImage<T>::Allocate(int res_x, int res_y)
{
	if(image_source.IsNotNull())
	{
		image_source = NULL;
	}

	image_source = ITKImageType::New();

	ITKImageType::RegionType region;
	ITKImageType::SizeType size;
	size[0] = res_x;
	size[1] = res_y;

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

}


template <typename T>
rts_itkImage<T> rts_itkImage<T>::SubImage(int px, int py, int sx, int sy)
{
	/*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(px >= DimX()) px = DimX() - 1;
	if(py >= DimY()) py = DimY() - 1;
	if(px + sx > DimX()) sx = DimX() - px - 1;
	if(py + sy > DimY()) sy = DimY() - py - 1;

	cout<<px<<","<<py<<","<<sx<<","<<sy<<endl;

	//create the resulting volume
	rts_itkImage<T> destination(sx, sy);

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

	//get an iterator for the source region
	ITKImageType::RegionType region;
	ITKImageType::IndexType index;
	ITKImageType::SizeType size;
	index[0] = px; index[1] = py;
	size[0] = sx; size[1] = sy;
	region.SetIndex(index);
	region.SetSize(size);
	itk::ImageRegionIterator<ITKImageType> s(image_source, 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<typename T>
T* rts_itkImage<T>::GetPointer()
{
	return image_source->GetBufferPointer();
}

template<typename T>
void rts_itkImage<T>::LoadImage(string filename)
{
	typedef itk::ImageFileReader<ITKImageType> ReaderType;
	ReaderType::Pointer reader = ReaderType::New();
	reader->SetFileName(filename.c_str());
	image_source = reader->GetOutput();

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

template <typename T>
void rts_itkImage<T>::LoadRAW(int header_size, int size_x,
							   int size_y, string filename)
{
	Allocate(size_x, size_y);


	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*)image_source->GetBufferPointer(), DimX()*DimY()*sizeof(T));
	
	infile.close();
}

template<typename T>
void rts_itkImage<T>::SaveImage(string filename)
{
	typedef itk::ImageFileWriter<ITKImageType> WriterType;
	WriterType::Pointer writer = WriterType::New();
	writer->SetFileName(filename.c_str());
	writer->SetInput(image_source);

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

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

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

template<typename T>
void rts_itkImage<T>::CastAndSaveImage(string filename)
{
	//create the destination image
	typedef itk::Image<unsigned char> DestinationType;
	DestinationType::Pointer dest_image = DestinationType::New();

	//create the casting filter
	typedef itk::CastImageFilter<ITKImageType, DestinationType> CastFilterType;
	CastFilterType::Pointer castFilter = CastFilterType::New();
	castFilter->SetInput(image_source);


	//cast the current image
	dest_image = castFilter->GetOutput();

	//save the result
	typedef itk::ImageFileWriter<DestinationType> WriterType;
	WriterType::Pointer writer = WriterType::New();
	writer->SetFileName(filename.c_str());
	writer->SetInput(dest_image);

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

template<typename T>
void rts_itkImage<T>::RescaleImage(T min_v, T max_v)
{
	typedef itk::RescaleIntensityImageFilter<ITKImageType, ITKImageType> RescaleFilter;
	RescaleFilter::Pointer rescaleFilter = RescaleFilter::New();
	rescaleFilter->SetInput(image_source);
	rescaleFilter->SetOutputMinimum(min_v);
	rescaleFilter->SetOutputMaximum(max_v);
	image_source = rescaleFilter->GetOutput();
	rescaleFilter->Update();
}

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

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

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

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

	//copy values from the DT Grid to the image
	ITKImageType::IndexType index;		//create an image index
	rtsDTGrid2D<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;
		image_source->SetPixel(index, i.Value());

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




}



#endif