/*********************************************************************************
* 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
********************************************************************************/
/*
Copyright (c) 2008 College of William and Mary
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

   * Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
   * Redistributions in binary form must reproduce the above
copyright notice, this list of conditions and the following disclaimer
in the documentation and/or other materials provided with the
distribution.
   * Neither the name of the College of William and Mary nor the
names of its contributors may be used to endorse or promote products
derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/


#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 = 2;

  typedef itk::Vector<float, 2> DisplacementType;
  typedef itk::Image<DisplacementType, 2> 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, 2> TransformType;
  typedef itk::Point<double, 2> 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;
  
  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];
 

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

  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){
 
        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));
        
        df->SetPixel(id, d);
        dfPixelCount++;

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

        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 < 9)
    {
    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;
}