/*=========================================================================
  This code provided as part of 
  Brooks, R., Arbel, T. "Improving the OrientedImage class for use in the
  Registration Framework", Insight Journal, October, 2007

  Author: Rupert Brooks
  Institution: Centre for Intelligent Machines, McGill University

  This example based closely on Examples/Registration/ImageRegistration3.cxx

=========================================================================*/
#if defined(_MSC_VER)
#pragma warning ( disable : 4786 )
#endif

#define DEG2RAD 3.14159265358979/180.0
#define TOL 1e-10

#include "itkImageRegistrationMethod.h"
#include "itkTranslationTransform.h"
#include "itkMeanSquaresImageToImageMetric.h"
#include "itkNormalizedCorrelationImageToImageMetric.h"
#include "itkMattesMutualInformationImageToImageMetric.h"
#include "itkFastOrientedImage.h"
#include "itkMeanSquaresModImageToImageMetric.h"
#include "itkNormalizedCorrelationModImageToImageMetric.h"
#include "itkMattesMutualInformationModImageToImageMetric.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkRegularStepGradientDescentOptimizer.h"
#include "itkSubtractImageFilter.h"
#include "itkRescaleIntensityImageFilter.h"
#include "itkOrientedImage.h"
#include "itkChangeInformationImageFilter.h"

#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"

#include "itkResampleImageFilter.h"
#include "itkCastImageFilter.h"


#include "itkCommand.h"

#include "time.h"
#include <ostream>
#include <fstream>

class CommandIterationUpdate : public itk::Command 
{
public:
  typedef  CommandIterationUpdate   Self;
  typedef  itk::Command             Superclass;
  typedef itk::SmartPointer<Self>  Pointer;
  std::ostream* outfile;
  itkNewMacro( Self );
protected:
	CommandIterationUpdate() {};

public:
  typedef itk::RegularStepGradientDescentOptimizer     OptimizerType;
  typedef const OptimizerType                         *OptimizerPointer;
  void Execute(itk::Object *caller, const itk::EventObject & event)
  {
    Execute( (const itk::Object *)caller, event);
  }
  void Execute(const itk::Object * object, const itk::EventObject & event)
  {
    OptimizerPointer optimizer = 
                         dynamic_cast< OptimizerPointer >( object );
    if( ! itk::IterationEvent().CheckEvent( &event ) )
      {
      return;
      }
      (*outfile) << optimizer->GetCurrentIteration() << " = ";
      (*outfile) << optimizer->GetValue() << " : ";
      (*outfile) << optimizer->GetCurrentPosition() << std::endl;
  }
   
};


int main( int argc, char *argv[] ){
// Test program for image metrics modified to correct derivative
// relative to the direction cosines of an image.
//
// Basically this test program operates by inputting an image, and
// creating a short filter chain that changes the direction cosines
//
// When the direction cosines are the identity matrix, the new
// classes should operate the same as the old classes.  When the 
// direction cosines are changed, the new classes should continue
// to work, while the old ones fail.
//
// As part of the modification, there is a choice of how to compute
// the derivative.  The old derivative is compared to the new one
// when the same method is being used, the result should be identical
// to the old method.

  if( argc < 4 )
    {
    std::cerr << "Missing Parameters " << std::endl;
    std::cerr << "Usage: " << argv[0];
    std::cerr << " flip metric derivativetype testImageFile [logfile]";
    std::cerr << "" << std::endl;
	std::cerr << "flip means:"<<std::endl; 
	std::cerr << "     0 - no rotation"<<std::endl; 
	std::cerr << "     1 - 90 degrees"<<std::endl; 
	std::cerr << "     2 - +30 degrees"<<std::endl; 
	std::cerr << "     3 - -30 degrees"<<std::endl; 
	std::cerr << "     4 - +120 degrees"<<std::endl; 
	std::cerr << "     5 - -120 degrees"<<std::endl; 
	std::cerr << "metric means:"<<std::endl; 
	std::cerr << "     0 - msd"<<std::endl; 
	std::cerr << "     1 - ncc"<<std::endl; 
	std::cerr << "     2 - mi"<<std::endl; 
	std::cerr << "derivative means:"<<std::endl; 
	std::cerr << "     0 - compute gradient image"<<std::endl; 
	std::cerr << "     1 - bspline"<<std::endl; 
	std::cerr << "     2 - derivative calculator"<<std::endl; 
    return 1;
    }
  int flip=1;
  int metricType=0;
  int derivativeType=0;

  std::ostream* outfile=NULL;
  std::ofstream logfile;
  if(argc>5) {
	  logfile.open(argv[5]);
	  outfile=&logfile;
  } else {
	  outfile=&std::cout;
  }


  const    unsigned int    Dimension = 2;
  typedef  unsigned short  PixelType;
  
  typedef itk::FastOrientedImage< PixelType, Dimension >  FixedImageType;
  typedef itk::FastOrientedImage< PixelType, Dimension >  MovingImageType;

  typedef itk::ChangeInformationImageFilter<FixedImageType> FixedChangeFilterType;
  typedef itk::ChangeInformationImageFilter<MovingImageType> MovingChangeFilterType;

  typedef itk::TranslationTransform< double, Dimension > TransformType;

  typedef itk::RegularStepGradientDescentOptimizer       OptimizerType;

  typedef itk::InterpolateImageFunction< 
                                    MovingImageType,
                                    double             > InterpolatorType;
  typedef itk::LinearInterpolateImageFunction< 
                                    MovingImageType,
                                    double             > InterpolatorType1;
  typedef itk::BSplineInterpolateImageFunction< 
                                    MovingImageType,
                                    double             > InterpolatorType2;

  typedef itk::ImageRegistrationMethod< 
                                    FixedImageType, 
                                    MovingImageType   >  RegistrationType;

  typedef itk::ImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType;
  typedef itk::MeanSquaresImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType1a;
  typedef itk::MeanSquaresModImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType2a;
  typedef itk::NormalizedCorrelationImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType1b;
  typedef itk::NormalizedCorrelationModImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType2b;
  typedef itk::MattesMutualInformationImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType1c;
  typedef itk::MattesMutualInformationModImageToImageMetric< 
                                      FixedImageType, 
                                      MovingImageType >  MetricType2c;

  typedef itk::ImageFileReader< FixedImageType  > FixedImageReaderType;
  typedef itk::ImageFileReader< MovingImageType > MovingImageReaderType;

  FixedImageReaderType::Pointer  fixedImageReader  = FixedImageReaderType::New();
  MovingImageReaderType::Pointer movingImageReader = MovingImageReaderType::New();
  flip=atoi(argv[1]);
  metricType=atoi(argv[2]);
  derivativeType=atoi(argv[3]);
  clock_t oldstart,oldend,newstart,newend;
  bool oldregistration,newregistration,derivsame;

  // flip the image, if desired
  FixedChangeFilterType::DirectionType directions;
  switch (flip) {
	case 0:
		(*outfile)<<"I'm orienting the images along the axes"<<std::endl;
		directions[0][0]=1.0;
		directions[0][1]=0.0;
		directions[1][0]=0.0;
		directions[1][1]=1.0;
		break;
	case 1:
		(*outfile)<<"I'm rotating the images 90 degrees"<<std::endl;
		directions[0][0]=0.0;
		directions[0][1]=-1.0;
		directions[1][0]=1.0;
		directions[1][1]=0.0;
		break;
	case 2:
		(*outfile)<<"I'm rotating the images +30 degrees"<<std::endl;
		directions[0][0]=cos(30.0*DEG2RAD);
		directions[0][1]=-sin(30.0*DEG2RAD);
		directions[1][0]=sin(30.0*DEG2RAD);
		directions[1][1]=cos(30.0*DEG2RAD);
		break;
	case 3:
		(*outfile)<<"I'm rotating the images -30 degrees"<<std::endl;
		directions[0][0]=cos(-30.0*DEG2RAD);
		directions[0][1]=-sin(-30.0*DEG2RAD);
		directions[1][0]=sin(-30.0*DEG2RAD);
		directions[1][1]=cos(-30.0*DEG2RAD);
		break;
	case 4:
		(*outfile)<<"I'm rotating the images +120 degrees"<<std::endl;
		directions[0][0]=cos(120.0*DEG2RAD);
		directions[0][1]=-sin(120.0*DEG2RAD);
		directions[1][0]=sin(120.0*DEG2RAD);
		directions[1][1]=cos(120.0*DEG2RAD);
		break;
	case 5:
		(*outfile)<<"I'm rotating the images -120 degrees"<<std::endl;
		directions[0][0]=cos(-120.0*DEG2RAD);
		directions[0][1]=-sin(-120.0*DEG2RAD);
		directions[1][0]=sin(-120.0*DEG2RAD);
		directions[1][1]=cos(-120.0*DEG2RAD);
		break;
	default:
		(*outfile)<<"Unknown flip mode "<<flip<<std::endl;
		exit(-1);
  }

  // read image and set up short filter chain
  fixedImageReader->SetFileName(  argv[4] );
  movingImageReader->SetFileName( argv[4] );

  FixedChangeFilterType::Pointer fixedChanger=FixedChangeFilterType::New();
  MovingChangeFilterType::Pointer movingChanger=MovingChangeFilterType::New();

  movingChanger->SetInput(movingImageReader->GetOutput());
  fixedChanger->SetInput(fixedImageReader->GetOutput());
  movingChanger->SetOutputDirection(directions);
  fixedChanger->SetOutputDirection(directions);
  fixedChanger->ChangeDirectionOn();
  movingChanger->ChangeDirectionOn();

  fixedChanger->Update(); // This is needed to make the BufferedRegion below valid.

  TransformType::Pointer      transform     = TransformType::New();
  MetricType::DerivativeType oldderivative(transform->GetNumberOfParameters());
  MetricType::DerivativeType newderivative(transform->GetNumberOfParameters());


  // each registration process is in its own block
  // they are nearly identical
  {
	OptimizerType::Pointer      optimizer     = OptimizerType::New();
	InterpolatorType::Pointer   interpolator;
	if(derivativeType==MetricType2a::BSpline) {
		(*outfile)<<"Using Bspline interpolator"<<std::endl;
		interpolator= InterpolatorType2::New();
	} else {
		interpolator= InterpolatorType1::New();
	}
	RegistrationType::Pointer   registration  = RegistrationType::New();

	registration->SetOptimizer(     optimizer     );
  
	MetricType::Pointer         metric;
	switch (metricType) {
		case 0:
			{
			MetricType1a::Pointer metric0=MetricType1a::New();
			metric=metric0;
			}
			break;
		case 1:
			{
			MetricType1b::Pointer metric0=MetricType1b::New();
			metric=metric0;
			}
			break;
		case 2:
			{
			MetricType1c::Pointer metric0=MetricType1c::New();
			metric0->ReinitializeSeed(0);
			metric0->SetNumberOfSpatialSamples(
				fixedChanger->GetOutput()->GetBufferedRegion().GetNumberOfPixels());
			metric=metric0;
			}
			break;
		default:
			std::cerr<<"Unknown metric "<<metricType<<std::endl;
	}
  
	registration->SetMetric( metric  );


	typedef RegistrationType::ParametersType ParametersType;
  
	ParametersType initialParameters=transform->GetParameters();
	initialParameters[0]=5.0 * directions[0][0]+3.0*directions[0][1];
	initialParameters[1]=5.0 * directions[1][0]+3.0*directions[1][1];

 
	metric->SetFixedImage(    fixedChanger->GetOutput()    );
	metric->SetMovingImage(   movingChanger->GetOutput()   );
	metric->SetFixedImageRegion( 
	fixedChanger->GetOutput()->GetBufferedRegion() );
	metric->SetTransform(     transform     );
	metric->SetInterpolator(  interpolator  );

	double value;
	try{
		metric->Initialize();
		metric->GetValueAndDerivative(initialParameters,value,oldderivative);
	} catch( itk::ExceptionObject & err ) { 
		(*outfile) << "ExceptionObject caught !" << std::endl; 
		(*outfile) << err << std::endl; 
		//return -1;
    } 
	
	(*outfile)<<"Probing metric at initial position:"<<std::endl;
	(*outfile)<<"Value: "<<value<<" derivative: "<<oldderivative<<std::endl;

	registration->SetInitialTransformParameters( initialParameters );
	registration->SetFixedImage(    fixedChanger->GetOutput()    );
	registration->SetMovingImage(   movingChanger->GetOutput()   );
	registration->SetTransform(     transform     );
	registration->SetInterpolator(  interpolator  );

	registration->SetFixedImageRegion( 
       fixedImageReader->GetOutput()->GetBufferedRegion() );


	optimizer->SetMaximumStepLength( 4.00 );  
	optimizer->SetMinimumStepLength( 0.01 );
	optimizer->SetNumberOfIterations( 200 );

	optimizer->MaximizeOff();


	CommandIterationUpdate::Pointer observer = CommandIterationUpdate::New();
	observer->outfile=outfile;
    optimizer->AddObserver( itk::IterationEvent(), observer );
	oldstart=clock();
	try { 
		registration->StartRegistration(); 
    } catch( itk::ExceptionObject & err ) { 
		(*outfile) << "ExceptionObject caught !" << std::endl; 
		(*outfile) << err << std::endl; 
		//return -1;
    } 
	oldend=clock();


	ParametersType finalParameters = registration->GetLastTransformParameters();
  
	const double TranslationAlongX = finalParameters[0];
	const double TranslationAlongY = finalParameters[1];
  
	oldregistration=(fabs(TranslationAlongX)<0.02 &&fabs(TranslationAlongY)<0.02);
	const unsigned int numberOfIterations = optimizer->GetCurrentIteration();
  
	const double bestValue = optimizer->GetValue();

	(*outfile) << "Registration done !" << std::endl;
	(*outfile) << "Number of iterations = " << numberOfIterations << std::endl;
	(*outfile) << "Parameters: "<<finalParameters<< std::endl;
	(*outfile) << "Optimal metric value = " << bestValue << std::endl;
	///8888888888888888888888888888888888
  }

  // now register with the new metric
  {
	OptimizerType::Pointer      optimizer     = OptimizerType::New();
	InterpolatorType::Pointer   interpolator;
	if(derivativeType==MetricType2a::BSpline) {
		interpolator= InterpolatorType2::New();
	} else {
		interpolator= InterpolatorType1::New();
	}
	RegistrationType::Pointer   registration  = RegistrationType::New();

	registration->SetOptimizer(     optimizer     );
  
	MetricType::Pointer         metric;
	switch (metricType) {
		case 0:
			{
			MetricType2a::Pointer metric0=MetricType2a::New();
			//metric0->SetNumberOfSpatialSamples(
			//fixedChanger->GetOutput()->GetBufferedRegion().GetNumberOfPixels());
			metric0->SetGradientMethod((MetricType2a::GradientCalculationMethod) derivativeType);
			metric=metric0;
			}
			break;
		case 1:
			{
			MetricType2b::Pointer metric0=MetricType2b::New();
			//metric0->SetNumberOfSpatialSamples(
			//	fixedChanger->GetOutput()->GetBufferedRegion().GetNumberOfPixels());
			metric0->SetGradientMethod((MetricType2b::GradientCalculationMethod)derivativeType);
			metric=metric0;
			}
			break;
		case 2:
			{
			MetricType2c::Pointer metric0=MetricType2c::New();
			metric0->SetNumberOfSpatialSamples(
				fixedChanger->GetOutput()->GetBufferedRegion().GetNumberOfPixels());
			metric0->ReinitializeSeed(0);
			metric0->SetGradientMethod((MetricType2c::GradientCalculationMethod)derivativeType);
			metric=metric0;
			}
			break;
		default:
			std::cerr<<"Unknown metric "<<metricType<<std::endl;
	}


  
	registration->SetMetric( metric  );


	typedef RegistrationType::ParametersType ParametersType;
  
	ParametersType initialParameters=transform->GetParameters();
	initialParameters[0]=5.0 * directions[0][0]+3.0*directions[0][1];
	initialParameters[1]=5.0 * directions[1][0]+3.0*directions[1][1];

 
	metric->SetFixedImage(    fixedChanger->GetOutput()    );
	metric->SetMovingImage(   movingChanger->GetOutput()   );
	metric->SetFixedImageRegion( 
	fixedChanger->GetOutput()->GetBufferedRegion() );
	metric->SetTransform(     transform     );
	metric->SetInterpolator(  interpolator  );

	double value;
	try{
		metric->Initialize();
		metric->GetValueAndDerivative(initialParameters,value,newderivative);
	} catch( itk::ExceptionObject & err ) { 
		(*outfile) << "ExceptionObject caught !" << std::endl; 
		(*outfile) << err << std::endl; 
		return -1;
    } 
	
	(*outfile)<<"Probing metric at initial position:"<<std::endl;
	(*outfile)<<"Value: "<<value<<" derivative: "<<newderivative<<std::endl;

	registration->SetInitialTransformParameters( initialParameters );
	registration->SetFixedImage(    fixedChanger->GetOutput()    );
	registration->SetMovingImage(   movingChanger->GetOutput()   );
	registration->SetTransform(     transform     );
	registration->SetInterpolator(  interpolator  );

	registration->SetFixedImageRegion( 
       fixedImageReader->GetOutput()->GetBufferedRegion() );


	optimizer->SetMaximumStepLength( 4.00 );  
	optimizer->SetMinimumStepLength( 0.01 );
	optimizer->SetNumberOfIterations( 200 );

	optimizer->MaximizeOff();


	CommandIterationUpdate::Pointer observer = CommandIterationUpdate::New();
	observer->outfile=outfile;
	optimizer->AddObserver( itk::IterationEvent(), observer );
	newstart=clock();
	try { 
		registration->StartRegistration(); 
    } catch( itk::ExceptionObject & err ) { 
		(*outfile) << "ExceptionObject caught !" << std::endl; 
		(*outfile) << err << std::endl; 
		return -1;
    } 
	newend=clock();


	ParametersType finalParameters = registration->GetLastTransformParameters();
  
	const double TranslationAlongX = finalParameters[0];
	const double TranslationAlongY = finalParameters[1];
	newregistration=(fabs(TranslationAlongX)<0.02 &&fabs(TranslationAlongY)<0.02);
  
	const unsigned int numberOfIterations = optimizer->GetCurrentIteration();
  
	const double bestValue = optimizer->GetValue();

	(*outfile) << "Registration done !" << std::endl;
	(*outfile) << "Number of iterations = " << numberOfIterations << std::endl;
	(*outfile) << "Parameters: "<<finalParameters<< std::endl;
	(*outfile) << "Optimal metric value = " << bestValue << std::endl;
	///8888888888888888888888888888888888
  }
  // compare derivatives
  derivsame=true;
  for(unsigned int i=0;i<transform->GetNumberOfParameters();i++) {
	  if(fabs(oldderivative[i]-newderivative[i]) > TOL ) derivsame=false;
  }
  if(derivsame){
	  (*outfile)<<"Derivative is the same"<<std::endl;
  } else {
	  (*outfile)<<"Derivative is not the same"<<std::endl;
  }
	  (*outfile)<<"Old Derivative: "<<oldderivative<<std::endl;
	  (*outfile)<<"New Derivative: "<<newderivative<<std::endl;
  if(oldregistration) {
	  (*outfile)<<"Registration with old method successful"<<std::endl;
  } else {
	  (*outfile)<<"Registration with old method failed"<<std::endl;
  }
  if(newregistration) {
	  (*outfile)<<"Registration with new method successful"<<std::endl;
  } else {
	  (*outfile)<<"Registration with new method failed"<<std::endl;
  }
	bool pass=newregistration &&(derivsame || flip ||
		((derivativeType>0&&metricType<2) ||(derivativeType==0 &&metricType==2)));
	// those last two things are the cases where its allowed to fail
  // timing differences
	  (*outfile)<<"Time Required (old method): "<<((double)(oldend-oldstart))/CLOCKS_PER_SEC<<std::endl;
	  (*outfile)<<"Time Required (new method): "<<((double)(newend-newstart))/CLOCKS_PER_SEC<<std::endl;
	  if(!pass) (*outfile)<<"ERROR: This run performed unexpectedly!"<<std::endl;
  return !pass;
}