/*   DTConverter
 read in the .dt DTI file and convert to nrrd or Nifti format
 ---Li Yang, May 2008
*/

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


int main(int argc, char *argv[])
{

  if (argc<9)
  {
      std::cout<<"Usage: DTConverter input Xsize Ysize Zsize spaceX spaceY spaceZ output\n";
	  std::cout<<"  input: the input DT image in local '.dt' format\n";
	  std::cout<<"  xsize ysize zsize: the size of the input DT image in X Y Z dimension\n";
	  std::cout<<"  spaceX spaceY spaceZ: the spacing of the input DT image in X Y Z dimension\n";
	  std::cout<<"  output: the output file name, either in Nrrd or Nifti format\n";
	  exit(0);
  }  
 
  //get the input image size and spacing
  int Xdim= atoi(argv[2]);
  int Ydim= atoi(argv[3]);
  int Zdim= atoi(argv[4]);

  float Xspa= atof(argv[5]);
  float Yspa= atof(argv[6]);
  float Zspa= atof(argv[7]);

  //read in the file--------------------------------
  FILE * pFile;
  long lSize;

  pFile = fopen ( argv[1] , "rb" );
  if (pFile==NULL) 
    return false;

  // obtain file size.
  fseek (pFile , 0 , SEEK_END);
  lSize = ftell (pFile);
  rewind (pFile);

  if ( lSize != 6*Xdim*Ydim*Zdim*sizeof(float))
  {
    fclose(pFile);
	std::cout<<"input file size is wrong!";
    return false;
  }

  //make a DT image	
  typedef itk::DiffusionTensor3D<float> PixelType;
  typedef itk::Image< PixelType, 3 > ImageType;

  ImageType::PixelType pixelValue;
  ImageType::IndexType pixelIndex;

  ImageType::Pointer image = ImageType::New();

  ImageType::IndexType start;
  start[0] = 0; // first index on X
  start[1] = 0; // first index on Y
  start[2] = 0; // first index on Z

  ImageType::SizeType size;
  size[0] = Xdim; // size along X
  size[1] = Ydim; // size along Y
  size[2] = Zdim; // size along Z

  ImageType::SpacingType spacing;
  spacing[0] = Xspa; // spacing along X
  spacing[1] = Yspa; // spacing along Y
  spacing[2] = Zspa; // spacing along Z
  image->SetSpacing(spacing);

  ImageType::RegionType region;
  region.SetSize( size );
  region.SetIndex( start );
  image->SetRegions( region );
  image->Allocate();

  float* element=new float;
  
  //write the nrhd or nii file
  typedef itk::ImageFileWriter< ImageType > WriterType;
  typedef itk::NrrdImageIO ImageIOTypeNrrd;
  typedef itk::NiftiImageIO ImageIOTypeNifti;

  WriterType::Pointer writer = WriterType::New();
  ImageIOTypeNrrd::Pointer NrrdIO = ImageIOTypeNrrd::New();
  ImageIOTypeNifti::Pointer NiftiIO = ImageIOTypeNifti::New();

  if ((strstr(argv[8],".nhdr")!=NULL)||(strstr(argv[8],".nrrd")!=NULL))
  {
	  writer->SetImageIO( NrrdIO );
	  for (int i=0; i<Zdim; i++)
	  {
		  for (int j=0; j<Ydim; j++)
		  {
			  for (int k=0; k<Xdim; k++)
			  {
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[0]=*element;	
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[3]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[5]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[1]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[2]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[4]=*element;

				  pixelIndex[0]=k;
				  pixelIndex[1]=j;
				  pixelIndex[2]=i;

				  image->SetPixel( pixelIndex, pixelValue );
			  }
		  }
	  } 
  }
  else if ((strstr(argv[8],".nii")!=NULL)||(strstr(argv[8],".hdr")!=NULL))
  {
	  writer->SetImageIO( NiftiIO );
	  for (int i=0; i<Zdim; i++)
	  {
		  for (int j=0; j<Ydim; j++)
		  {
			  for (int k=0; k<Xdim; k++)
			  {
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[0]=*element;	
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[2]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[5]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[1]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[3]=*element;
				  fread(element, sizeof(float), 1, pFile);
				  pixelValue[4]=*element;

				  pixelIndex[0]=k;
				  pixelIndex[1]=j;
				  pixelIndex[2]=i;

				  image->SetPixel( pixelIndex, pixelValue );
			  }
		  }
	  } 
  } 
  else
  {
	  std::cout<<"only support nrrd or Nifti format\n";
	  fclose (pFile);
	  exit(0);
  }  
  
  fclose (pFile);
  writer->SetFileName( argv[8] );
  writer->SetInput( image );
  
  try
  {
	writer->Update();
  }
  catch( itk::ExceptionObject & err )
  {
	std::cerr << "ExceptionObject caught !" << std::endl;
	std::cerr << err << std::endl;
	return EXIT_FAILURE;
  }
}