/*********************************************************************************
* DFGenerator.cxx
*
* This tool takes an input image and a binary mask of the object of interest, and
* creates a synthetic deformation field (DF) for subsequent registration 
* and/or similarity metric evaluation.
*
*
*Authors: Andriy Fedorov, Eric Billet 
*PI: Nikos Chrisochoides
*
*********************************************************************************/


#include "itkExtractImageFilter.h"
#include "itkImage.h"
#include "itkImageRegionIteratorWithIndex.h"
#include "itkImageRegionIterator.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkRegionOfInterestImageFilter.h"
#include "itkDanielssonDistanceMapImageFilter.h"
#include "itkNumericTraits.h"
//#include <vector>
#include "itkVector.h"
#include <time.h>
#include "itkImageConstIterator.h"
#include "itkOrientImageFilter.h"
#include "itkMersenneTwisterRandomVariateGenerator.h"
#include "itkGaussianDistribution.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkThinPlateSplineKernelTransform.h"
#include "itkResampleImageFilter.h"
#include "itkPointSet.h"
#include "itkMaskImageFilter.h"
#include <iostream>
#include "itkPasteImageFilter.h"
#include "itkCastImageFilter.h"
#include "itkExtractImageFilter.h"
#include "itkNearestNeighborInterpolateImageFunction.h"

#define MT_SEED 123


const unsigned int Dimension = 3;

  typedef itk::Vector<float, 3> DisplacementType;
  typedef itk::Image<DisplacementType, 3> DFImageType;
  typedef itk::ImageFileReader< DFImageType  >  DFReaderType;
  typedef float        InputPixelType;
  typedef itk::Image< InputPixelType,  Dimension >    InputImageType;
  typedef itk::ImageFileReader< InputImageType  >  ReaderType;
  typedef itk::ImageFileWriter< InputImageType >  WriterType;
  typedef itk::ImageFileWriter<DFImageType> DFWriterType;

 
  typedef itk::ImageRegionIterator<InputImageType> IteratorType;
  typedef itk::ImageRegionIterator<DFImageType> DFIteratorType;
  typedef itk::Statistics::GaussianDistribution DistributionType;
  typedef itk::Statistics::MersenneTwisterRandomVariateGenerator UniformDistributionType;

  typedef itk::ThinPlateSplineKernelTransform<double, 3> TransformType;
  typedef itk::Point<double, 3> PointType;
  typedef TransformType::PointSetType PointSetType;
  typedef PointSetType::PointIdentifier PointIdType;
  typedef std::vector<PointType> PointArrayType;
  typedef itk::ResampleImageFilter<InputImageType, InputImageType> ResamplerType;
  typedef itk::LinearInterpolateImageFunction<InputImageType, double> InterpolatorType;
  typedef itk::NearestNeighborInterpolateImageFunction<InputImageType, double> MaskInterpolatorType;
  typedef itk::OrientImageFilter<InputImageType, InputImageType> ImageOrienter;
  typedef itk::OrientImageFilter<DFImageType, DFImageType> DFOrienter;


/**
This function will deform the input image by generating displacement vectors at
sparse locations defined by setting the sampling spacing.  Each vector is then
drawn from a Gaussian distribution.  The dense deformation field is evaluated
at each point of the image by Thin Plate Splines interpolation. 
*/
void DeformImage(char * image1,char * mask,float df_spacing,float var, InputImageType::Pointer &deformed_image, DFImageType::Pointer &df, InputImageType::Pointer &deformed_mask){
	
  ReaderType::Pointer maskReader = ReaderType::New();
  ReaderType::Pointer imageReader = ReaderType::New();
  imageReader->SetFileName(image1);
  maskReader->SetFileName(mask);
  imageReader->Update();
  maskReader->Update();

  InputImageType::Pointer inputImage = imageReader->GetOutput();
  InputImageType::Pointer inputMask = maskReader->GetOutput();

  DFImageType::RegionType dfRegion;
  DFImageType::PixelType zeroDispl;
  zeroDispl[0] = 0;
  zeroDispl[1] = 0;
  zeroDispl[2] = 0;
  dfRegion.SetSize(inputImage->GetLargestPossibleRegion().GetSize());
  dfRegion.SetIndex(inputImage->GetLargestPossibleRegion().GetIndex());
  df->SetRegions(dfRegion);
  df->SetSpacing(inputImage->GetSpacing());
  df->Allocate();
  df->FillBuffer(zeroDispl);


  double xMax, yMax, zMax, x0, y0, z0;

  xMax = dfRegion.GetSize()[0]*inputImage->GetSpacing()[0];
  yMax = dfRegion.GetSize()[1]*inputImage->GetSpacing()[1];
  zMax = dfRegion.GetSize()[2]*inputImage->GetSpacing()[2];

  x0 = dfRegion.GetIndex()[0]*inputImage->GetSpacing()[0];
  y0 = dfRegion.GetIndex()[1]*inputImage->GetSpacing()[1];
  z0 = dfRegion.GetIndex()[2]*inputImage->GetSpacing()[2];

  UniformDistributionType::Pointer uniformDistr = UniformDistributionType::New();
  uniformDistr->Initialize(MT_SEED);

  DistributionType::Pointer distr = DistributionType::New();
  distr->SetMean(0);
  distr->SetVariance(var);

  InputImageType::PointType pt;
  InputImageType::IndexType id;
  unsigned dfPixelCount = 0;
  DFImageType::PixelType d;

  PointSetType::Pointer srcLM = PointSetType::New();
  PointSetType::Pointer tgtLM = PointSetType::New();

  PointSetType::PointsContainer::Pointer srcLMPts = srcLM->GetPoints();
  PointSetType::PointsContainer::Pointer tgtLMPts = tgtLM->GetPoints();

  PointIdType pid = itk::NumericTraits<PointIdType>::Zero;

  // Generate Gaussian deformation field
  for(pt[0]=x0;pt[0]<xMax;pt[0]+=df_spacing){
    for(pt[1]=y0;pt[1]<yMax;pt[1]+=df_spacing){
      for(pt[2]=z0;pt[2]<zMax;pt[2]+=df_spacing){
        if(!inputMask->TransformPhysicalPointToIndex(pt, id))
          continue; // point outside the image
        if(!inputMask->GetPixel(id))
          continue; // image mask zero at index
        d[0] = distr->EvaluateInverseCDF(uniformDistr->GetUniformVariate(0.0, 1.0));
        d[1] = distr->EvaluateInverseCDF(uniformDistr->GetUniformVariate(0.0, 1.0));
        d[2] = distr->EvaluateInverseCDF(uniformDistr->GetUniformVariate(0.0, 1.0));
        df->SetPixel(id, d);
        dfPixelCount++;

        PointType sLM, tLM;
        sLM[0] = pt[0];
        sLM[1] = pt[1];
        sLM[2] = pt[2];
        
        tLM[0] = pt[0]+d[0];
        tLM[1] = pt[1]+d[1];
        tLM[2] = pt[2]+d[2];

        srcLM->SetPoint(pid, sLM);
        tgtLM->SetPoint(pid, tLM);
        pid++;
      }
    }
  }

  

  // TPS interpolation of the translated points
  TransformType::Pointer transform = TransformType::New();
  transform->SetSourceLandmarks(srcLM.GetPointer());
  transform->SetTargetLandmarks(tgtLM.GetPointer());
  transform->ComputeWMatrix();

  

  DFIteratorType dfI(df, df->GetBufferedRegion());
  DFImageType::IndexType dfIndex;
  DFImageType::PointType dfPoint, dfResamplPoint;
  DFImageType::PixelType dfDisplacement;
  for(dfI.GoToBegin();!dfI.IsAtEnd();++dfI){
    dfIndex = dfI.GetIndex();
    if(!inputImage->GetPixel(dfIndex))
      continue; // image mask zero at index
    df->TransformIndexToPhysicalPoint(dfIndex, dfPoint);
    dfResamplPoint = transform->TransformPoint(dfPoint);
    // NB: assume displacement doesn't go out of image boundaries
    dfDisplacement = dfPoint-dfResamplPoint;
    dfI.Set(dfDisplacement);
  }

  
  ResamplerType::Pointer resampler = ResamplerType::New();
  InterpolatorType::Pointer interpolator = InterpolatorType::New();
  resampler->SetInterpolator( interpolator );
  InputImageType::PointType   origin  = inputImage->GetOrigin();
  InputImageType::RegionType region = inputImage->GetBufferedRegion();
  InputImageType::SizeType   size =  region.GetSize();

  resampler->SetOutputSpacing(inputImage->GetSpacing());
  resampler->SetOutputOrigin(origin);
  resampler->SetSize(size);
  resampler->SetTransform(transform);
  resampler->SetOutputStartIndex(region.GetIndex());
  resampler->SetInput(inputImage);

  //Deform the mask
  ResamplerType::Pointer resampler_mask = ResamplerType::New();
  MaskInterpolatorType::Pointer interpolator_mask = MaskInterpolatorType::New();
  resampler_mask->SetInterpolator( interpolator_mask );
  InputImageType::PointType   origin_mask  = inputMask->GetOrigin();
  InputImageType::RegionType region_mask = inputMask->GetBufferedRegion();
  InputImageType::SizeType   size_mask =  region_mask.GetSize();

  resampler_mask->SetOutputSpacing(inputMask->GetSpacing());
  resampler_mask->SetOutputOrigin(origin_mask);
  resampler_mask->SetSize(size_mask);
  resampler_mask->SetTransform(transform);
  resampler_mask->SetOutputStartIndex(region_mask.GetIndex());
  resampler_mask->SetInput(inputMask);
  //write deformed mask
  resampler_mask->Update();
  deformed_mask = resampler_mask->GetOutput();


  
  //save output to deformed_image
  resampler->Update();
  deformed_image = resampler->GetOutput();

}



int main( int argc, char ** argv )
{
	
  // Verify the number of parameters in the command line
  if( argc < 1 )
    {
    std::cerr << "Usage: " << std::endl;
    std::cerr << argv[0] << "[image][mask][variance][df_spacing][output_image][output_df][output_mask][output_dfnorm]" << std::endl;
    return EXIT_FAILURE;
    }

	

	double maximum, RMS_error;
	char * moving_image_name=argv[1];
	char * mask_name = argv[2];
	float var = atof(argv[3]);
	float df_spacing = atof(argv[4]);
	
	//Deform the image
	DFImageType::Pointer df = DFImageType::New();
	InputImageType::Pointer deformed_image= InputImageType::New();
	InputImageType::Pointer deformed_mask= InputImageType::New();
	DeformImage(moving_image_name,mask_name,df_spacing,var, deformed_image, df, deformed_mask);



        //generate deformation field norm image for error comparison
        InputImageType::Pointer DFmapImage = InputImageType::New();
	DFmapImage->SetRegions( deformed_image->GetRequestedRegion() );
	DFmapImage->SetSpacing(deformed_image->GetSpacing());
	DFmapImage->Allocate();


	DFIteratorType DFIt(df, df->GetRequestedRegion());
	IteratorType  DFmapIt(DFmapImage, df->GetRequestedRegion());
	

	//iterate through images
	for(DFIt.GoToBegin(), DFmapIt.GoToBegin() ; !DFIt.IsAtEnd(); ++DFIt, ++DFmapIt){
		//find norm
		double dist = 0;
		DFImageType::PixelType df;
		df = DFIt.Get();
		DFmapIt.Set(df.GetNorm());
         
	}



         //write out the images
        WriterType::Pointer writerDF = WriterType::New();
        writerDF->SetFileName(argv[8]);
        writerDF->SetInput( DFmapImage );
        writerDF->Update();



	WriterType::Pointer writer2 = WriterType::New();
	writer2->SetFileName(argv[5]);
	writer2->SetInput( deformed_image );
    	writer2->Update();
	WriterType::Pointer writer3 = WriterType::New();
	writer3->SetFileName(argv[7]);
	writer3->SetInput( deformed_mask );
    	writer3->Update();
	DFWriterType::Pointer writerdf = DFWriterType::New();
	writerdf->SetFileName(argv[6]);
	writerdf->SetInput( df );
    	writerdf->Update();
	


return 0;
}