/*=========================================================================
 *
 *  Copyright Insight Software Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#include "itkTranslationTransform.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkNormalizedCorrelationPointSetToImageMetric.h"
#include "itkRegularStepGradientDescentOptimizer.h"
#include "itkPointSet.h"
#include "itkPointSetToImageRegistrationMethod.h"
#include "itkImageRegistrationMethodImageSource.h"

#include <iostream>

/**
 *
 *  This program tests the registration of a PointSet against an image.
 *
 *  This test uses two 2D-Gaussians (standard deviation RegionSize/2)
 *  One is shifted by 5 pixels from the other.
 *
 */

int itkPointSetToImageRegistrationTest_1(int, char* [] )
{

//------------------------------------------------------------
// Create two simple images
//------------------------------------------------------------

  const unsigned int ImageDimension = 2;

  typedef double                   PixelType;

  typedef double                   CoordinateRepresentationType;

  //Allocate Images
  typedef itk::Image<PixelType,ImageDimension>   MovingImageType;
  typedef itk::Image<PixelType,ImageDimension>   FixedImageType;

  // ImageSource
  typedef itk::testhelper::ImageRegistrationMethodImageSource<
    PixelType,
    PixelType,
    ImageDimension >   ImageSourceType;

  ImageSourceType::Pointer    imageSource   = ImageSourceType::New();

  itk::Size<ImageDimension> size;
  size[0] = 100;
  size[1] = 100;

  imageSource->GenerateImages( size );

  MovingImageType::ConstPointer movingImage = imageSource->GetMovingImage();
  FixedImageType::ConstPointer  fixedImage  = imageSource->GetFixedImage();

//-----------------------------------------------------------
// Create the point set and load it with data by sampling
// the fixed image
//-----------------------------------------------------------
  typedef itk::PointSet< float, 2 >   FixedPointSetType;
  FixedPointSetType::Pointer fixedPointSet = FixedPointSetType::New();

  const unsigned int numberOfPoints = 10000;

  fixedPointSet->SetPointData( FixedPointSetType::PointDataContainer::New() );

  fixedPointSet->GetPoints()->Reserve( numberOfPoints );
  fixedPointSet->GetPointData()->Reserve( numberOfPoints );

  typedef itk::ImageRegionConstIterator< FixedImageType > ImageIteratorType;

  ImageIteratorType it( fixedImage, fixedImage->GetBufferedRegion() );

  const unsigned int skip =
    fixedImage->GetBufferedRegion().GetNumberOfPixels() / numberOfPoints;

  unsigned int counter = 0;

  FixedPointSetType::PointIdentifier pointId = 0;
  FixedPointSetType::PointType  point;

  it.GoToBegin();
  while( !it.IsAtEnd() )
    {
    if( counter==0 )
      {
      fixedImage->TransformIndexToPhysicalPoint( it.GetIndex(), point );
      fixedPointSet->SetPoint( pointId, point );
      fixedPointSet->SetPointData( pointId, it.Get() );
      ++pointId;
      if( pointId == numberOfPoints )
        {
        break;
        }
      counter = skip;
      }
    --counter;
    ++it;
    }

//-----------------------------------------------------------
// Set up  the Metric
//-----------------------------------------------------------
  typedef itk::NormalizedCorrelationPointSetToImageMetric<
    FixedPointSetType,
    MovingImageType >
    MetricType;

  typedef MetricType::TransformType                 TransformBaseType;
  typedef TransformBaseType::ParametersType         ParametersType;

  MetricType::Pointer  metric = MetricType::New();


//-----------------------------------------------------------
// Set up a Transform
//-----------------------------------------------------------

  typedef itk::TranslationTransform<
    CoordinateRepresentationType,
    ImageDimension >         TransformType;

  TransformType::Pointer transform = TransformType::New();


//------------------------------------------------------------
// Set up an Interpolator
//------------------------------------------------------------
  typedef itk::LinearInterpolateImageFunction<
    MovingImageType,
    double > InterpolatorType;

  InterpolatorType::Pointer interpolator = InterpolatorType::New();

  // Optimizer Type
  typedef itk::RegularStepGradientDescentOptimizer       OptimizerType;

  OptimizerType::Pointer      optimizer     = OptimizerType::New();

  // Registration Method
  typedef itk::PointSetToImageRegistrationMethod<
    FixedPointSetType,
    MovingImageType >    RegistrationType;


  RegistrationType::Pointer   registration  = RegistrationType::New();


  typedef itk::CommandIterationUpdate<
    OptimizerType >    CommandIterationType;


  // Instantiate an Observer to report the progress of the Optimization
  CommandIterationType::Pointer iterationCommand = CommandIterationType::New();
  iterationCommand->SetOptimizer(  optimizer.GetPointer() );


  // Scale the translation components of the Transform in the Optimizer
  OptimizerType::ScalesType scales( transform->GetNumberOfParameters() );
  scales.Fill( 1.0 );


  unsigned long   numberOfIterations =   50;
  //double          translationScale   =  1.0; // not used
  double          maximumStepLenght  =  1.0;  // no step will be larger than this
  double          minimumStepLenght  =  0.01; // convergence criterion
  double          gradientTolerance  =  1e-6; // convergence criterion


  optimizer->SetScales( scales );
  optimizer->SetNumberOfIterations( numberOfIterations );
  optimizer->SetMinimumStepLength( minimumStepLenght );
  optimizer->SetMaximumStepLength( maximumStepLenght );
  optimizer->SetGradientMagnitudeTolerance( gradientTolerance );
  optimizer->MinimizeOn();

  // Start from an Identity transform (in a normal case, the user
  // can probably provide a better guess than the identity...
  transform->SetIdentity();
  registration->SetInitialTransformParameters( transform->GetParameters() );

  //------------------------------------------------------
  // Connect all the components required for Registration
  //------------------------------------------------------
  registration->SetMetric(        metric        );
  registration->SetOptimizer(     optimizer     );
  registration->SetTransform(     transform     );
  registration->SetFixedPointSet( fixedPointSet );
  registration->SetMovingImage(   movingImage   );
  registration->SetInterpolator(  interpolator  );

//------------------------------------------------------------
// Set up transform parameters
//------------------------------------------------------------
  ParametersType parameters( transform->GetNumberOfParameters() );

  // initialize the offset/vector part
  for( unsigned int k = 0; k < ImageDimension; k++ )
    {
    parameters[k]= 0.0f;
    }


  try
    {
    registration->Update();
    }
  catch( itk::ExceptionObject & e )
    {
    std::cout << e << std::endl;
    return EXIT_FAILURE;
    }

  return EXIT_SUCCESS;

}
