/*
 * Copyright (c) ICG. All rights reserved.
 * See copyright.txt for more information.
 *
 * Institute for Computer Graphics and Vision
 * Graz, University of Technology / Austria
 *
 *
 * This software is distributed WITHOUT ANY WARRANTY; without even
 * the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
 * PURPOSE.  See the above copyright notices for more information.
 *
 *
 * Project     : MIPItkProjects
 * Module      : Evaluation
 * Class       : $RCSfile: ApplySyntheticTPSTransform.cxx,v $
 * Language    : C++
 * Description : 
 *
 * Author     : Martin Urschler
 * EMail      : urschler@icg.tu-graz.ac.at
 * Date       : $Date: 2007-03-07 10:01:10 $
 * Version    : $Revision: 1.8 $
 * Full Id    : $Id: ApplySyntheticTPSTransform.cxx,v 1.8 2007-03-07 10:01:10 urschler Exp $
 *
 */

#include "commonItkTypedefs.h"
#include "SmallUtilityMethods.h"
#include "VolumeIOWrapper.h"
#include "itkThinPlateSplineKernelTransform.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkImageFileWriter.h"
#include "vnl/vnl_random.h"




int main( int argc, char * argv[] )
{
   if( argc != 6 )
   {
       std::cerr << "Missing Parameters " << std::endl;
       std::cerr << "Usage: " << argv[0];
       std::cerr << "  fixedImage  warpedImage  gridSize  maxDeviation" <<
          "  resampleDefaultValue" << std::endl;
       return EXIT_FAILURE;
   }
    
   std::string inputFilename( argv[1] );
   std::string outputFilename( argv[2] );
   const int gridSize = atoi( argv[3] );
   const int maxDeviation = vnl_math_min( gridSize, atoi(argv[4]) );
   std::cout << "maxDeviation: " << maxDeviation << std::endl;
   const Int16ImageType::PixelType resampleDefaultValue = atoi( argv[5] );
   
   
   Int16ImageType::Pointer originalImage =
      VolumeIOWrapper<Int16ImageType>::readITKVolume( inputFilename, "" );

   Int16ImageType::SizeType originalImageSize = 
      originalImage->GetLargestPossibleRegion().GetSize();

   typedef itk::ThinPlateSplineKernelTransform<double> KernelTransformType;
   KernelTransformType::Pointer kernelTransform = KernelTransformType::New();
   
   typedef KernelTransformType::PointSetType SplineTransformPointSetType;
   typedef SplineTransformPointSetType::PointType SplineTransformPointType;

   SplineTransformPointSetType::Pointer sourceLandmarks = 
      SplineTransformPointSetType::New();
   SplineTransformPointSetType::Pointer targetLandmarks = 
      SplineTransformPointSetType::New();

   const unsigned int sizeX = originalImageSize[0];
   const unsigned int sizeY = originalImageSize[1];
   const unsigned int sizeZ = originalImageSize[2];

   std::cout << "sizeX: " << sizeX << std::endl;
   std::cout << "sizeY: " << sizeY << std::endl;
   std::cout << "sizeZ: " << sizeZ << std::endl;


   // setup the grid and calculate the synthetic deformations
   {
      vnl_random randomNumberGenerator;
      
      unsigned long pointId = 0;
      
      Int16ImageType::IndexType sourceIndex, targetIndex;
      Int16ImageType::PointType sourcePoint, targetPoint;
      SplineTransformPointType splineSourcePoint, splineTargetPoint;
      
      for (int z=gridSize; z<sizeZ; z+=gridSize)
      {
         for (int y=gridSize; y<sizeY; y+=gridSize)
         {
            for (int x=gridSize; x<sizeX; x+=gridSize)
            {
               sourceIndex[0] = x;
               sourceIndex[1] = y;
               sourceIndex[2] = z;
               originalImage->TransformIndexToPhysicalPoint( sourceIndex, 
                  sourcePoint );
               
               splineSourcePoint[0] = sourcePoint[0];
               splineSourcePoint[1] = sourcePoint[1];
               splineSourcePoint[2] = sourcePoint[2];
               
               int x_displacement = 
                  randomNumberGenerator.lrand32(2*maxDeviation);
               x_displacement -= maxDeviation;
               int y_displacement = 
                  randomNumberGenerator.lrand32(2*maxDeviation);
               y_displacement -= maxDeviation;
               int z_displacement = 
                  randomNumberGenerator.lrand32(2*maxDeviation);
               z_displacement -= maxDeviation;
               
               targetIndex[0] = x+x_displacement;
               targetIndex[1] = y+y_displacement;
               targetIndex[2] = z+z_displacement;
               originalImage->TransformIndexToPhysicalPoint( targetIndex, 
                  targetPoint );
               
               splineTargetPoint[0] = targetPoint[0];
               splineTargetPoint[1] = targetPoint[1];
               splineTargetPoint[2] = targetPoint[2];
               
               //std::cout << splineSourcePoint << " -> " << splineTargetPoint 
               //   << std::endl; 
               
               sourceLandmarks->SetPoint( pointId, splineSourcePoint );
               targetLandmarks->SetPoint( pointId, splineTargetPoint );
               pointId++;
            }
         }
      }
      
      std::cout << "number of landmarks: " << pointId << std::endl;
   }
   

   kernelTransform->SetSourceLandmarks( sourceLandmarks.GetPointer() );
   kernelTransform->SetTargetLandmarks( targetLandmarks.GetPointer() );
   

   std::cout << "compute w matrix!" << std::endl;
   kernelTransform->ComputeWMatrix();

   std::cout << "warping!" << std::endl;
   
   typedef itk::LinearInterpolateImageFunction<Int16ImageType, double>
      InterpolatorType;
   InterpolatorType::Pointer interpolator = InterpolatorType::New();
   interpolator->SetInputImage(originalImage);
   
   Int16ImageType::Pointer outputImage = 
      SmallUtilityMethods<Int16ImageType>::createNewImage( originalImage );
   itk::ImageRegionIteratorWithIndex<Int16ImageType> outIter(
      outputImage, outputImage->GetLargestPossibleRegion());
   
   // create a result deformation field
   DeformationFieldType::Pointer deformationField = 
      SmallUtilityMethods<Int16ImageType>::createNewZeroedDeformationField( 
         outputImage );
   
   itk::ImageRegionIterator<DeformationFieldType> fieldIter( deformationField, 
      deformationField->GetLargestPossibleRegion() );
   
   
   Int16ImageType::IndexType outputIndex;
   Int16ImageType::PointType outputPoint;
   Int16ImageType::PointType inputResampleLocation;
   VectorPixelType deformation;
   
   while ( !outIter.IsAtEnd() )
   {
      outputIndex = outIter.GetIndex();
      outputImage->TransformIndexToPhysicalPoint( outputIndex, 
         outputPoint );
      
      // Compute corresponding displaced point 
      inputResampleLocation = kernelTransform->TransformPoint( outputPoint );
      
      if (interpolator->IsInsideBuffer(inputResampleLocation))
      {
         outIter.Set( ROUND_DBL<Int16ImageType::PixelType>(
            interpolator->Evaluate(inputResampleLocation)) );
            
         VectorPixelType deformation = inputResampleLocation - outputPoint;
         fieldIter.Set( deformation );
      }
      else
      {
         outIter.Set( resampleDefaultValue );
         fieldIter.Set( 
            VectorPixelConstants<VectorPixelType::ValueType, 3>::undefined() );
      }
      
      ++outIter;
      ++fieldIter;
   }
   
   // file output
   
   VolumeIOWrapper<Int16ImageType>::writeITKVolume16Bit( 
      outputImage, outputFilename, "" );
   
   
   std::string displacementFieldFilename = outputFilename;
   displacementFieldFilename.replace(
      displacementFieldFilename.find(".hdr"), 4, ".mha" ); 
   typedef itk::ImageFileWriter<DeformationFieldType> DeformationFieldWriter;
   DeformationFieldWriter::Pointer writer = DeformationFieldWriter::New();
   writer->SetFileName( displacementFieldFilename.c_str() );
   writer->SetInput( deformationField );
   std::cout << "Writing deformation field" << 
      displacementFieldFilename.c_str() << std::endl;
   writer->Update();

   return EXIT_SUCCESS;
}