/*   DTIPad
read in the DTI file (in nrrd or nifti format) and pad 0 around it and than save as .dt format
---Li Yang, May 2008
*/


#include <iostream>
#include <string.h>
#include <stdio.h>
#include "itkDiffusionTensor3D.h"
#include "itkImage.h"
#include "itkImageFileWriter.h"
#include "itkImageFileReader.h"
#include "itkNrrdImageIO.h"
#include "itkTensorFractionalAnisotropyImageFilter.h"
#include "itkNiftiImageIO.h"

int main(int argc, char *argv[])
{
	if(argc<9 )
	{
		std::cout<<"Usage: DTIPad input Xsize Ysize Zsize Xstart Ystart Zstart output\n";
		std::cout<<"  input: the input DT image in Nrrd or Nifti format\n";
		std::cout<<"  Xsize Ysize Zsize: the size of the padded output in .dt format\n";
		std::cout<<"  Xstart Ystart Zstart: position of the first voxel of the input in the padded image (from 1)\n";
		std::cout<<"  output: the output .dt image\n";
		exit(0);
	}  

	int *OutSize = new int[3];
	OutSize[0]=atoi(argv[2]);
	OutSize[1]=atoi(argv[3]);
	OutSize[2]=atoi(argv[4]);

	int *start = new int[3];
	start[0]=atoi(argv[5]);
	start[1]=atoi(argv[6]);
	start[2]=atoi(argv[7]);

	//get the input image
	typedef itk::DiffusionTensor3D<float> PixelType;
	typedef itk::Image< PixelType, 3 > ImageType;
	typedef itk::NrrdImageIO ImageIOTypeNrrd;
	typedef itk::NiftiImageIO ImageIOTypeNifti;
	typedef itk::ImageFileWriter< ImageType > WriterType;
	typedef itk::ImageFileReader< ImageType > ReaderType;

	ImageType::Pointer DTImage = ImageType::New();
	WriterType::Pointer writer = WriterType::New();
	ReaderType::Pointer reader = ReaderType::New();
	ImageIOTypeNrrd::Pointer NrrdIO = ImageIOTypeNrrd::New();
	ImageIOTypeNifti::Pointer NiftiIO = ImageIOTypeNifti::New();

	if ((strstr(argv[1],".nhdr")!=NULL)||(strstr(argv[1],".nrrd")!=NULL))
	{
		reader->SetImageIO( NrrdIO );
	}
	else if ((strstr(argv[1],".nii")!=NULL)||(strstr(argv[1],".hdr")!=NULL))
	{
		reader->SetImageIO( NiftiIO );
	} 
	else
	{
		std::cout<<"only support Nrrd or Nifti format\n";
		exit(0);
	}  	

	reader->SetFileName( argv[1] );
	reader->Update();
	DTImage = reader->GetOutput();

	//get the input image information, size spacing
	ImageType::RegionType Region = DTImage->GetLargestPossibleRegion();
	ImageType::SizeType Size = Region.GetSize();
	if ((OutSize[0]<=Size[0])||(OutSize[1]<=Size[1])||(OutSize[2]<=Size[2]))
	{
		std::cout<<"wrong image size\n";
		exit(0);
	}

	//swap pixelValue[2] and pixelValue[3], because FA filter is written for nrrd sequence only
	ImageType::PixelType pixelValue;
	ImageType::IndexType pixelIndex;

	if ((strstr(argv[1],".nii")!=NULL)||(strstr(argv[1],".hdr")!=NULL))
	{
		
		float temp = 0;

		for (int i=0; i<Size[0]; i++)
		{
			for (int j=0; j<Size[1]; j++)
			{
				for (int k=0; k<Size[2]; k++)
				{
					pixelIndex[0] = i;
					pixelIndex[1] = j;
					pixelIndex[2] = k;
					pixelValue = DTImage->GetPixel( pixelIndex );

					temp = pixelValue[2];
					pixelValue[2] = pixelValue[3];
					pixelValue[3] = temp;

					DTImage->SetPixel(pixelIndex, pixelValue);
				}
			}
		}
	} 

	//write out the .dt file
	FILE * pFile;
	pFile = fopen ( argv[8] , "wb" );

	float pad [] = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
	float *element = new float;

	for (int i=1;i<=OutSize[2];i++)
	{
		for (int j=1;j<=OutSize[1];j++)
		{
			for (int k=1;k<=OutSize[0];k++)
			{
				if ((k<start[0])||(k>=start[0]+Size[0])||(j<start[1])||(j>=start[1]+Size[1])||(i<start[2])||(i>=start[2]+Size[2]))
				{
					fwrite ( pad, sizeof(float), 6 , pFile );
				}
				else
				{
					pixelIndex[2] = i-start[2];
					pixelIndex[1] = j-start[1];
					pixelIndex[0] = k-start[0];					
					pixelValue = DTImage->GetPixel( pixelIndex );
					
					*element = pixelValue[0];
					fwrite ( element, sizeof(float), 1 , pFile );
					*element = pixelValue[3];
					fwrite ( element, sizeof(float), 1 , pFile );
					*element = pixelValue[5];
					fwrite ( element, sizeof(float), 1 , pFile );
					*element = pixelValue[1];
					fwrite ( element, sizeof(float), 1 , pFile );
					*element = pixelValue[2];
					fwrite ( element, sizeof(float), 1 , pFile );
					*element = pixelValue[4];
					fwrite ( element, sizeof(float), 1 , pFile );
				}
			}			
		}		
	}
	
	fclose (pFile);
	return 0;	
}