#include "itkImage.h"
#include "itkBSplineDeformableTransform.h"
#include "itkGridScheduleComputer.h"
#include "itkUpsampleBSplineParametersFilter.h"
#include "itkAffineTransform.h"
#include "vnl/vnl_math.h"
#include <iostream>


int main( int argc, char **argv )
{
  bool testPassed = true;

  /** Constants */
  const unsigned int splineOrder = 3;
  const unsigned int dimension   = 2;

  /** Typedefs */
  typedef itk::Image<
    unsigned char, dimension >                TestImageType;
  typedef TestImageType::IndexType            IndexType;
  typedef TestImageType::PointType            PointType;
  typedef double                              CoordRepType;
  typedef itk::BSplineDeformableTransform<
    CoordRepType, dimension, splineOrder >    TransformType;
  typedef itk::AffineTransform<
    CoordRepType, dimension >                 InitialTransformType;
  typedef 
    InitialTransformType::OutputVectorType    ScalesType;
  typedef TransformType::ParametersType       ParametersType;
  typedef ParametersType::ValueType           ParametersValueType;
  typedef TransformType::ImageType            CoefficientImageType;
  typedef itk::GridScheduleComputer<
    CoordRepType, dimension >                 ScheduleComputerType;
  typedef ScheduleComputerType::OriginType    OriginType;
  typedef ScheduleComputerType::SpacingType   SpacingType;
  typedef ScheduleComputerType::SizeType      SizeType;
  typedef ScheduleComputerType::RegionType    RegionType;
  typedef ScheduleComputerType::
    VectorGridSpacingFactorType               ScheduleType;
  typedef 
    itk::UpsampleBSplineParametersFilter<
    ParametersType, CoefficientImageType >    UpsamplerType;

  /** Instantiate some variables */
  TransformType::Pointer lowTransform = 
    TransformType::New();
  TransformType::Pointer highTransform =
    TransformType::New();
  InitialTransformType::Pointer initialTransform = 
    InitialTransformType::New();
  ScheduleComputerType::Pointer scheduleComputer =
    ScheduleComputerType::New();
  UpsamplerType::Pointer upsampler = 
    UpsamplerType::New();

  /** Properties of the (imaginary) image on which we want to 
   * overlay a B-spline control point grid. */
  OriginType imageOrigin;
  SpacingType imageSpacing;
  SizeType imageSize;
  RegionType imageRegion;
  imageOrigin.Fill(0.0);
  for ( unsigned int d = 0; d < dimension; ++d )
  {
    imageSpacing[d] = static_cast<double>( d+1 );
  }  
  imageSize.Fill(32);
  imageRegion.SetSize( imageSize );
  TestImageType::Pointer testImage = TestImageType::New();
  testImage->SetRegions( imageRegion );
  testImage->SetOrigin( imageOrigin );
  testImage->SetSpacing( imageSpacing );  
  IndexType imageEndIndex;
  for (unsigned int i = 0; i < dimension; ++i )
  {
    imageEndIndex[i] = imageSize[i]-1;
  }
  PointType imageEndPoint;
  testImage->TransformIndexToPhysicalPoint( imageEndIndex, imageEndPoint );
  
  /** The desired grid spacing in mm at the finest resolution level. */
  SpacingType finalGridSpacing;
  double finalGridSpacingIsotropic = 9.0;
  finalGridSpacing.Fill( finalGridSpacingIsotropic );
  
  /** The number of resolution levels */
  const unsigned int nrOfResolutions = 2;

  std::cout << "--- Testing class " << scheduleComputer->GetNameOfClass() 
    << " ---" << std::endl;

  /** Configure the scheduleComputer */
  scheduleComputer->SetImageOrigin( imageOrigin );
  scheduleComputer->SetImageSpacing( imageSpacing );
  scheduleComputer->SetImageRegion( imageRegion );
  scheduleComputer->SetBSplineOrder( splineOrder );
  scheduleComputer->SetFinalGridSpacing( finalGridSpacing );

  /** Test the SetDefaultSchedule method */
  std::cout << "\nTesting the SetDefaultSchedule() method." << std::endl;
  const double upsamplingFactor = 2.0;
  ScheduleType expectedSchedule;
  expectedSchedule.resize( nrOfResolutions );
  expectedSchedule[0].Fill( upsamplingFactor );
  expectedSchedule[1].Fill( 1.0 );  
  scheduleComputer->SetDefaultSchedule( nrOfResolutions, upsamplingFactor );
  ScheduleType schedule;  
  scheduleComputer->GetSchedule(schedule);

  std::cout << "Expected default schedule:\n";
  for ( unsigned int s = 0; s < expectedSchedule.size(); ++s )
  {
    std::cout << expectedSchedule[s] << std::endl;
  }
  std::cout << "Default schedule returned by "
    << scheduleComputer->GetNameOfClass() << ":" << std::endl;
  for ( unsigned int s = 0; s < schedule.size(); ++s )
  {
    std::cout << schedule[s] << std::endl;
  }
    
  bool defaultScheduleTestPassed = true;
  if ( schedule == expectedSchedule )
  {
    std::cout 
      << "The default schedule of the "
      << scheduleComputer->GetNameOfClass() 
      << " is correct." << std::endl;
  }
  else
  {
    defaultScheduleTestPassed = false;
    
    std::cout 
      << "The default schedule of the "
      << scheduleComputer->GetNameOfClass()
      << " is incorrect." << std::endl;    
  }
  testPassed = defaultScheduleTestPassed;
  
  /** Modify the schedule a bit and set it in the schedule computer */
  schedule = expectedSchedule;
  schedule[0][0] = 3.0;
  scheduleComputer->SetSchedule( schedule );

  /** Test the ComputeBSplineGrid method */
  std::cout << "\nTesting the ComputeBSplineGrid() method." << std::endl;
  std::cout << "Definition of the input image size:\n";
  std::cout << "ImageSize = " << imageRegion.GetSize() << "\n";
  std::cout << "ImageSpacing = " << imageSpacing << "\n";
  std::cout << "ImageOrigin = " << imageOrigin << "\n";
  std::cout << "Definition of the schedule and final grid spacing:\n";
  std::cout << "Schedule = \n";
  for (unsigned int s = 0; s < schedule.size(); ++s )
  {
    std::cout << schedule[s] << "\n";
  }
  std::cout << "FinalGridSpacing = " << scheduleComputer->GetFinalGridSpacing() << std::endl;
  
  scheduleComputer->ComputeBSplineGrid();
  
  RegionType gridRegion;
  SpacingType gridSpacing;
  OriginType gridOrigin;  
  CoefficientImageType::Pointer gridImage = CoefficientImageType::New();
  bool computeBSplineGridTestPassed = true;
  
  std::cout << "Result of the ComputeBSplineGrid() method:\n";
  for ( unsigned int s = 0; s < nrOfResolutions; ++s )
  {
    scheduleComputer->GetBSplineGrid( s, gridRegion, gridSpacing, gridOrigin );
    std::cout << "BSpline grid at resolution level " << s << ":\n";
    std::cout << "GridSize = " << gridRegion.GetSize() << "\n";
    std::cout << "GridSpacing = " << gridSpacing << "\n";
    std::cout << "GridOrigin = " << gridOrigin << "\n";

    /** Create an image according to the specified grid */
    gridImage->SetRegions( gridRegion );
    gridImage->SetSpacing( gridSpacing );
    gridImage->SetOrigin( gridOrigin );
    IndexType gridEndIndex;
    for (unsigned int i = 0; i < dimension; ++i )
    {
      gridEndIndex[i] = gridRegion.GetSize()[i] - 1;
    }
    PointType gridEndPoint;
    gridImage->TransformIndexToPhysicalPoint( gridEndIndex, gridEndPoint );  
        
    /** Do some checks on the grid. */
    for ( unsigned int d = 0; d < dimension; ++d )
    {
      /** The grid should be layed symmetric over the image. 
       * dif should equal 0. */
      const double dif = (imageOrigin[d] - gridOrigin[d]) - (gridEndPoint[d] - imageEndPoint[d]);
      computeBSplineGridTestPassed &= ( dif <= itk::NumericTraits<double>::epsilon() );      
      
      /** Left from the image origin should be at least (splineOrder/2 + 1) points.
       * dif2 should be greater than or equal to 0. */
      const double dif2 = ( imageOrigin[d] - gridOrigin[d] ) - ( ( splineOrder/2 ) * gridSpacing[d] );
      computeBSplineGridTestPassed &= ( dif2 >= itk::NumericTraits<double>::Zero );
    }

  }

  /** The final grid spacing should be equal to the desired final grid spacing */
  computeBSplineGridTestPassed &= (gridSpacing == finalGridSpacing);

  if ( computeBSplineGridTestPassed )
  {
    std::cout << "The ComputeBSplineGrid() test passed." << std::endl;
  }
  else
  {
    std::cout << "The ComputeBSplineGrid() test failed." << std::endl;
  }
  testPassed &= computeBSplineGridTestPassed;

  /** Test the SetInitialTransform() with an affine transform with isotropic scaling.*/
  std::cout << "\nTesting the SetInitialTransform() method." << std::endl;
  initialTransform->SetCenter( imageOrigin );
  initialTransform->SetIdentity();
  ScalesType scale;
  const double scaling = 3.0;
  scale.Fill(scaling);
  initialTransform->Scale(scale, true );
  scheduleComputer->SetInitialTransform( initialTransform );
  scheduleComputer->ComputeBSplineGrid();
  /** Check at the finest resolution level if the resulting grid spacing
   * is indeed 'scaling' times larger */
  scheduleComputer->GetBSplineGrid( nrOfResolutions - 1, gridRegion, gridSpacing, gridOrigin );
  bool initialTransformPassed = true;
  for ( unsigned int d = 0; d < dimension; ++d )
  {
    const double dif = vnl_math_abs( gridSpacing[d] / finalGridSpacing[d] / scaling - 1.0 );
    initialTransformPassed &= ( dif <= itk::NumericTraits<double>::epsilon() );    
  }

  if ( initialTransformPassed )
  {
    std::cout << "The SetInitialTransform() test passed." << std::endl;
  }
  else
  {
    std::cout << "The SetInitialTransform() test failed." << std::endl;
  }
  testPassed &= initialTransformPassed;

  /** Test the UpsampleBSplineParametersFilter. */
  std::cout << "\n--- Testing class " << upsampler->GetNameOfClass() 
    << " ---" << std::endl;

  /** Use a default schedule: */
  scheduleComputer->SetInitialTransform( 0 );
  scheduleComputer->SetDefaultSchedule( nrOfResolutions, upsamplingFactor );
  scheduleComputer->ComputeBSplineGrid(); 
  RegionType lowGridRegion;
  SpacingType lowGridSpacing;
  OriginType lowGridOrigin;
  RegionType highGridRegion;
  SpacingType highGridSpacing;
  OriginType highGridOrigin;  
  scheduleComputer->GetBSplineGrid( 0, lowGridRegion, lowGridSpacing, lowGridOrigin );
  scheduleComputer->GetBSplineGrid( 1, highGridRegion, highGridSpacing, highGridOrigin );

  /** Configure the upsampler */
  upsampler->SetCurrentGridOrigin( lowGridOrigin );  
  upsampler->SetCurrentGridSpacing( lowGridSpacing );
  upsampler->SetCurrentGridRegion( lowGridRegion );
  upsampler->SetBSplineOrder( splineOrder );
  upsampler->SetRequiredGridOrigin( highGridOrigin );
  upsampler->SetRequiredGridSpacing( highGridSpacing );
  upsampler->SetRequiredGridRegion( highGridRegion );

  /** Set up the bspline transforms */
  lowTransform->SetGridRegion( lowGridRegion );
  lowTransform->SetGridSpacing( lowGridSpacing );
  lowTransform->SetGridOrigin( lowGridOrigin );
  highTransform->SetGridRegion( highGridRegion );
  highTransform->SetGridSpacing( highGridSpacing );
  highTransform->SetGridOrigin( highGridOrigin );
  
  /** Do a trivial test: check whether zero parameters returns
   * zero parameters with the right array size: */
  std::cout << "\nTest if UpsampleParameters() keeps zero zero." << std::endl;
  ParametersType lowParameters( lowTransform->GetNumberOfParameters() ); 
  lowParameters.Fill(0.0);
  lowTransform->SetParameters( lowParameters );
  ParametersType highParameters;
  upsampler->UpsampleParameters( lowParameters, highParameters );
  bool zeroTestPassed = true;
  for ( unsigned int p = 0; p < highParameters.GetSize(); ++p )
  {
    zeroTestPassed &= ( vnl_math_abs( highParameters[p] ) <= 
      itk::NumericTraits<double>::epsilon() );
  }
  zeroTestPassed &= ( highTransform->GetNumberOfParameters() == highParameters.GetSize() );
  if ( zeroTestPassed )
  {
    std::cout << "Test if UpsampleParameters() keeps zero zero passed." << std::endl;
  }
  else
  {
    std::cout << "Test if UpsampleParameters() keeps zero zero failed." << std::endl;
  }
  testPassed &= zeroTestPassed;
  
  /** Reproducability test. See whether the results are the same as
   * they once were on a specific platform. */
  std::cout << "\nTest if the results are reproducable." << std::endl; 
  for ( unsigned int p = 0; p < lowParameters.GetSize(); ++p )
  {
    lowParameters[p] = vcl_cos( static_cast<double>( p ) );
  }
  upsampler->UpsampleParameters( lowParameters, highParameters );
  const unsigned int nrOfParamsToCheck = 5;
  itk::Array<unsigned int> checkIndices( nrOfParamsToCheck );
  ParametersType checkParameters( nrOfParamsToCheck );
  ParametersType testParameters( nrOfParamsToCheck );
  checkIndices[0] = 0;
  checkIndices[1] = 10;
  checkIndices[2] = 11;
  checkIndices[3] = 18;
  checkIndices[4] = highParameters.GetSize() - 2 ;
  checkParameters[0] = 0.80599539;
  checkParameters[1] = 0.60352076;
  checkParameters[2] = 0.29263144;
  checkParameters[3] = 0.56380516;
  checkParameters[4] = 0.55443502;

  bool reproTestPassed = true;
  for ( unsigned int i = 0; i < nrOfParamsToCheck; ++i )
  {
    testParameters[i] = highParameters[ checkIndices[i] ];
    reproTestPassed &= ( ( testParameters[i] - checkParameters[i] ) <= 1e-8 );
  }
  std::cout << std::showpoint;
  std::cout.precision( 8 );
  std::cout << "Checking parameter indices: " << checkIndices << std::endl;
  std::cout << "Expected values: " << checkParameters << std::endl;
  std::cout << "Measured values: " << testParameters << std::endl;

  if ( reproTestPassed )
  {
    std::cout << "Reproducability test passed." << std::endl;
  }
  else
  {
    std::cout << "Reproducability test failed." << std::endl;
  }
  testPassed &= reproTestPassed;

  /** Summarise test results. */  
  if ( testPassed )
  {
    std::cout << "\nAll tests in BSplineHelpersTest passed." << std::endl;
  }
  else
  {
    std::cout << "\nSome tests in BSplineHelpersTest failed." << std::endl;
  }

  /** Return 0 if test passed. 1 otherwise. */
  return static_cast<int>( !testPassed );

} // end main