/*
* Copyright (c) ICG. All rights reserved.
*
* 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     : projects
* Module      : MIPViewer - plugins
* Class       : $RCSfile: CTDiffeomorphicDemonsRegistration.txx,v $
* Language    : C++
* Description : 
*
* Author     : Martin Urschler
* EMail      : urschler@icg.tu-graz.ac.at
* Date       : $Date: 2007-07-11 15:34:20 $
* Version    : $Revision: 1.6 $
* Full Id    : $Id: CTDiffeomorphicDemonsRegistration.txx,v 1.6 2007-07-11 15:34:20 urschler Exp $
*
*/


#ifndef _CTDIFFEOMORPHICDEMONSREGISTRATION_TXX_
#define _CTDIFFEOMORPHICDEMONSREGISTRATION_TXX_


#include "VolumeIOWrapper.h"
#include "SmallUtilityMethods.h"
#include "NonlinearRegistrationUtilities.h"

#include "itkLinearInterpolateImageFunction.h"

#include "insight-journal/vercauteren/itkMultiResolutionPDEDeformableRegistration2.h"
#include "insight-journal/vercauteren/itkFastSymmetricForcesDemonsRegistrationFilter.h"
#include "insight-journal/vercauteren/itkDiffeomorphicDemonsRegistrationFilter.h"
#include "itkWarpImageFilter.h"
#include "itkCommand.h"
#include "insight-journal/vercauteren/itkWarpJacobianDeterminantFilter.h"
#include "itkMinimumMaximumImageCalculator.h"
#include "insight-journal/vercauteren/itkWarpSmoothnessCalculator.h"
#include "insight-journal/vercauteren/itkGridForwardWarpImageFilter.h"
#include "insight-journal/vercauteren/itkVectorCentralDifferenceImageFunction.h"

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




//  The following section of code implements a Command observer
//  that will monitor the evolution of the registration process.
//
template <class TPixel=float, unsigned int VImageDimension=3>
class CommandIterationUpdate : public itk::Command 
{
public:
   typedef  CommandIterationUpdate   Self;
   typedef  itk::Command             Superclass;
   typedef  itk::SmartPointer<Self>  Pointer;

   typedef itk::Image< TPixel, VImageDimension > InternalImageType;
   typedef itk::Vector< TPixel, VImageDimension >    VectorPixelType;
   typedef itk::Image<  VectorPixelType, VImageDimension > DeformationFieldType;

   typedef itk::DiffeomorphicDemonsRegistrationFilter<
      InternalImageType,
      InternalImageType,
      DeformationFieldType>   DiffeomorphicDemonsRegistrationFilterType;

   typedef itk::FastSymmetricForcesDemonsRegistrationFilter<
      InternalImageType,
      InternalImageType,
      DeformationFieldType>   FastSymmetricForcesDemonsRegistrationFilterType;

   typedef itk::MultiResolutionPDEDeformableRegistration2<
      InternalImageType, InternalImageType,
      DeformationFieldType, TPixel >   MultiResRegistrationFilterType;

   typedef itk::WarpJacobianDeterminantFilter<
      DeformationFieldType, InternalImageType> JacobianFilterType;
   
   typedef itk::MinimumMaximumImageCalculator<InternalImageType> MinMaxFilterType;

   typedef itk::WarpSmoothnessCalculator<DeformationFieldType>
      SmoothnessCalculatorType;

   typedef itk::VectorCentralDifferenceImageFunction<DeformationFieldType>
      WarpGradientCalculatorType;

   typedef typename WarpGradientCalculatorType::OutputType WarpGradientType;
   
   itkNewMacro( Self );

private:
   std::ofstream m_Fid;
   bool m_headerwritten;
   typename JacobianFilterType::Pointer m_JacobianFilter;
   typename MinMaxFilterType::Pointer m_Minmaxfilter;
   typename SmoothnessCalculatorType::Pointer m_SmothnessCalculator;
   typename DeformationFieldType::ConstPointer m_TrueField;
   typename WarpGradientCalculatorType::Pointer m_TrueWarpGradientCalculator;
   typename WarpGradientCalculatorType::Pointer m_CompWarpGradientCalculator;

public:
   void SetTrueField(const DeformationFieldType * truefield)
   {
      m_TrueField = truefield;

      m_TrueWarpGradientCalculator = WarpGradientCalculatorType::New();
      m_TrueWarpGradientCalculator->SetInputImage( m_TrueField );

      m_CompWarpGradientCalculator =  WarpGradientCalculatorType::New();
   }
   
   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)
   {
      if( !(itk::IterationEvent().CheckEvent( &event )) )
      {
         return;
      }

      typename DeformationFieldType::ConstPointer deffield = 0;
      unsigned int iter = -1;
      double metricbefore = -1.0;
      
      if ( const DiffeomorphicDemonsRegistrationFilterType * filter = 
           dynamic_cast< const DiffeomorphicDemonsRegistrationFilterType * >( object ) )
      {
         iter = filter->GetElapsedIterations() - 1;
         metricbefore = filter->GetMetric();
         deffield = const_cast<DiffeomorphicDemonsRegistrationFilterType *>
            (filter)->GetDeformationField();
      }
      else if ( const FastSymmetricForcesDemonsRegistrationFilterType * filter = 
           dynamic_cast< const FastSymmetricForcesDemonsRegistrationFilterType * >( object ) )
      {
         iter = filter->GetElapsedIterations() - 1;
         metricbefore = filter->GetMetric();
         deffield = const_cast<FastSymmetricForcesDemonsRegistrationFilterType *>
            (filter)->GetDeformationField();
      }
      else if ( const MultiResRegistrationFilterType * multiresfilter = 
           dynamic_cast< const MultiResRegistrationFilterType * >( object ) )
      {
         std::cout<<"Finished Multi-resolution iteration :"<<multiresfilter->GetCurrentLevel()-1<<std::endl;
         std::cout<<"=============================="<<std::endl<<std::endl;
      }
      else
      {
         return;
      }

      if (deffield)
      {
         std::cout<<iter<<": MSE "<<metricbefore<<" - ";

         double fieldDist = -1.0;
         double fieldGradDist = -1.0;
         double tmp;
         if (m_TrueField)
         {
            typedef itk::ImageRegionConstIteratorWithIndex<DeformationFieldType>
               FieldIteratorType;
            FieldIteratorType currIter(
               deffield, deffield->GetLargestPossibleRegion() );
            FieldIteratorType trueIter(
               m_TrueField, deffield->GetLargestPossibleRegion() );

            m_CompWarpGradientCalculator->SetInputImage( deffield );

            fieldDist = 0.0;
            fieldGradDist = 0.0;
            for ( currIter.GoToBegin(), trueIter.GoToBegin();
                  !currIter.IsAtEnd(); ++currIter, ++trueIter )
            {
               fieldDist += (currIter.Value() - trueIter.Value()).GetSquaredNorm();

               // No need to add Id matrix here as we do a substraction
               tmp = (
                  ( m_CompWarpGradientCalculator->EvaluateAtIndex(currIter.GetIndex())
                    -m_TrueWarpGradientCalculator->EvaluateAtIndex(trueIter.GetIndex())
                     ).GetVnlMatrix() ).frobenius_norm();
               fieldGradDist += tmp*tmp;
            }
            fieldDist = sqrt( fieldDist/ (double)(
                     deffield->GetLargestPossibleRegion().GetNumberOfPixels()) );
            fieldGradDist = sqrt( fieldGradDist/ (double)(
                     deffield->GetLargestPossibleRegion().GetNumberOfPixels()) );
            
            std::cout<<"d(.,true) "<<fieldDist<<" - ";
            std::cout<<"d(.,Jac(true)) "<<fieldGradDist<<" - ";
         }
         
         m_SmothnessCalculator->SetImage( deffield );
         m_SmothnessCalculator->Compute();
         const double harmonicEnergy = m_SmothnessCalculator->GetSmoothness();
         std::cout<<"harmo. "<<harmonicEnergy<<" - ";

         
         m_JacobianFilter->SetInput( deffield );
         m_JacobianFilter->UpdateLargestPossibleRegion();

        
         const unsigned int numPix = m_JacobianFilter->
            GetOutput()->GetLargestPossibleRegion().GetNumberOfPixels();
         
         TPixel* pix_start = m_JacobianFilter->GetOutput()->GetBufferPointer();
         TPixel* pix_end = pix_start + numPix;

         TPixel* jac_ptr;

         // Get percentage of det(Jac) below 0
         unsigned int jacBelowZero(0u);
         for (jac_ptr=pix_start; jac_ptr!=pix_end; ++jac_ptr)
         {
            if ( *jac_ptr<=0.0 ) ++jacBelowZero;
         }
         const double jacBelowZeroPrc = static_cast<double>(jacBelowZero)
            / static_cast<double>(numPix);
         

         // Get min an max jac
         /*
         std::pair<TPixel*, TPixel*> minmax_res =
            boost::minmax_element(pix_start, pix_end);
         */

         //const double minJac = *(minmax_res.first);
         //const double maxJac = *(minmax_res.second);

         const double minJac = *(std::min_element (pix_start, pix_end));
         const double maxJac = *(std::max_element (pix_start, pix_end));

         // Get some quantiles
         // We don't need the jacobian image
         // we can modify/sort it in place
         jac_ptr = pix_start + static_cast<unsigned int>(0.002*numPix);
         std::nth_element(pix_start, jac_ptr, pix_end);
         const double Q002 = *jac_ptr;

         jac_ptr = pix_start + static_cast<unsigned int>(0.01*numPix);
         std::nth_element(pix_start, jac_ptr, pix_end);
         const double Q01 = *jac_ptr;

         jac_ptr = pix_start + static_cast<unsigned int>(0.99*numPix);
         std::nth_element(pix_start, jac_ptr, pix_end);
         const double Q99 = *jac_ptr;

         jac_ptr = pix_start + static_cast<unsigned int>(0.998*numPix);
         std::nth_element(pix_start, jac_ptr, pix_end);
         const double Q998 = *jac_ptr;
         

         std::cout<<"max|Jac| "<<maxJac<<" - "
                  <<"min|Jac| "<<minJac<<" - "
                  <<"ratio(|Jac|<=0) "<<jacBelowZeroPrc<<std::endl;
         
         

         if (this->m_Fid.is_open())
         {
            if (!m_headerwritten)
            {
               this->m_Fid<<"Iteration"
                          <<", MSE before"
                          <<", Harmonic energy"
                          <<", min|Jac|"
                          <<", 0.2% |Jac|"
                          <<", 01% |Jac|"
                          <<", 99% |Jac|"
                          <<", 99.8% |Jac|"
                          <<", max|Jac|"
                          <<", ratio(|Jac|<=0)";
               
               if (m_TrueField)
               {
                  this->m_Fid<<", dist(warp,true warp)"
                             <<", dist(Jac,true Jac)";
               }
               
               this->m_Fid<<std::endl;
               
               m_headerwritten = true;
            }
            
            this->m_Fid<<iter
                       <<", "<<metricbefore
                       <<", "<<harmonicEnergy
                       <<", "<<minJac
                       <<", "<<Q002
                       <<", "<<Q01
                       <<", "<<Q99
                       <<", "<<Q998
                       <<", "<<maxJac
                       <<", "<<jacBelowZeroPrc;

            if (m_TrueField)
            {
               this->m_Fid<<", "<<fieldDist
                          <<", "<<fieldGradDist;
            }
            
            this->m_Fid<<std::endl;
         }
      }
   }
   
protected:   
   CommandIterationUpdate() :
      m_Fid( "metricvalues.csv" ),
      m_headerwritten(false)
   {
      m_JacobianFilter = JacobianFilterType::New();
      m_JacobianFilter->SetUseImageSpacing( true );
      m_JacobianFilter->ReleaseDataFlagOn();
      
      m_Minmaxfilter = MinMaxFilterType::New();

      m_SmothnessCalculator = SmoothnessCalculatorType::New();

      m_TrueField = 0;
      m_TrueWarpGradientCalculator = 0;
      m_CompWarpGradientCalculator = 0;
   };

   ~CommandIterationUpdate()
   {
      this->m_Fid.close();
   }
};



template< typename TPixelDataType >
CTNonlinearRegistration::ReturnType 
CTDiffeomorphicDemonsRegistration<TPixelDataType>::
matching(
   const unsigned int maxNrShrinkLevels, 
   const unsigned int iterationsParam1, 
   const unsigned int iterationsParam2,
   const float sigmaDeformationField,
   const float sigmaUpdateField,
   const float maxLengthUpdateField,
   const unsigned int demonsForceGradientType )
{
   const bool useDiffeomorphicDemons = true;
   const int verbosityMode = 1;
   
   ITK_IMAGE_ALLOCATE_MACRO(FloatImageType3D, fixedImage, 
      InputImageType, this->m_fixedImage);
   ITK_IMAGE_COPY_MACRO(InputImageType, this->m_fixedImage, 
      FloatImageType3D, fixedImage);
   ITK_IMAGE_ALLOCATE_MACRO(FloatImageType3D, movingImage, 
      InputImageType, this->m_movingImage);
   ITK_IMAGE_COPY_MACRO(InputImageType, this->m_movingImage, 
      FloatImageType3D, movingImage);
      

   // Set up the demons filter output
   typename DeformationFieldType::Pointer defField = 0;

   {//for mem allocations
   
   // Set up the demons filter
   typedef typename itk::PDEDeformableRegistrationFilter
      < FloatImageType3D, FloatImageType3D, DeformationFieldType>
      BaseRegistrationFilterType;
   typename BaseRegistrationFilterType::Pointer filter;
   
   if ( useDiffeomorphicDemons == true )
   {
      typedef typename itk::DiffeomorphicDemonsRegistrationFilter
         < FloatImageType3D, FloatImageType3D, DeformationFieldType>   ActualRegistrationFilterType;
      typedef typename ActualRegistrationFilterType::GradientType GradientType;
      
      typename ActualRegistrationFilterType::Pointer actualfilter = ActualRegistrationFilterType::New();

      actualfilter->SetMaximumUpdateStepLength( maxLengthUpdateField );
      actualfilter->SetUseGradientType( static_cast<GradientType>(demonsForceGradientType) );
      filter = actualfilter;
   }
   else
   {
      typedef typename itk::FastSymmetricForcesDemonsRegistrationFilter
         < FloatImageType3D, FloatImageType3D, DeformationFieldType>   ActualRegistrationFilterType;
      typedef typename ActualRegistrationFilterType::GradientType GradientType;
      
      typename ActualRegistrationFilterType::Pointer actualfilter = ActualRegistrationFilterType::New();
      
      actualfilter->SetMaximumUpdateStepLength( maxLengthUpdateField );
      actualfilter->SetUseGradientType( static_cast<GradientType>(demonsForceGradientType) );
      filter = actualfilter;
   }

   if ( sigmaDeformationField > 0.1 )
   {
      filter->SmoothDeformationFieldOn();
      filter->SetStandardDeviations( sigmaDeformationField );
   }
   else
      filter->SmoothDeformationFieldOff();

   if ( sigmaUpdateField > 0.1 )
   {
      filter->SmoothUpdateFieldOn();
      filter->SetUpdateFieldStandardDeviations( sigmaUpdateField );
   }
   else
      filter->SmoothUpdateFieldOff();

   //filter->SetIntensityDifferenceThreshold( 0.001 );

   if ( verbosityMode > 0 )
   {      
      // Create the Command observer and register it with the registration filter.
      typename CommandIterationUpdate<float, 3>::Pointer observer =
         CommandIterationUpdate<float, 3>::New();

      
      filter->AddObserver( itk::IterationEvent(), observer );
   }

   // Set up the multi-resolution filter
   typedef typename itk::MultiResolutionPDEDeformableRegistration2<
      FloatImageType3D, FloatImageType3D, DeformationFieldType, float >   MultiResRegistrationFilterType;
   typename MultiResRegistrationFilterType::Pointer multires = MultiResRegistrationFilterType::New();

   multires->SetRegistrationFilter( filter );
   multires->SetNumberOfLevels( maxNrShrinkLevels );
   
   std::vector<unsigned int> iters;
   for (unsigned int i=0; i<maxNrShrinkLevels; ++i)
   {
      iters.push_back(iterationsParam1 * (maxNrShrinkLevels - i - 1) + iterationsParam2);
   }
   
   multires->SetNumberOfIterations( &iters[0] );

   multires->SetFixedImage( fixedImage );
   multires->SetMovingImage( movingImage );


   if ( verbosityMode > 0 )
   {
      // Create the Command observer and register it with the registration filter.
      typename CommandIterationUpdate<float, 3>::Pointer multiresobserver =
         CommandIterationUpdate<float, 3>::New();
      multires->AddObserver( itk::IterationEvent(), multiresobserver );
   }
   

   
   // Compute the deformation field
   try
   {
      multires->UpdateLargestPossibleRegion();
   }
   catch( itk::ExceptionObject& err )
   {
      std::cout << "Unexpected error." << std::endl;
      std::cout << err << std::endl;
      exit( EXIT_FAILURE );
   }


   // The outputs
   defField = multires->GetOutput();
   defField->DisconnectPipeline();

   }//end for mem allocations

   
   // warp the result
   typedef itk::WarpImageFilter
      < FloatImageType3D, FloatImageType3D, DeformationFieldType >  WarperType;
   typename WarperType::Pointer warper = WarperType::New();
   warper->SetInput( movingImage );
   warper->SetOutputSpacing( fixedImage->GetSpacing() );
   warper->SetOutputOrigin( fixedImage->GetOrigin() );
   warper->SetDeformationField( defField );
   warper->SetEdgePaddingValue( this->m_resampleDefaultValue );

   EXCEPTION_CHECKED_ITK_FILTER_UPDATE( "final image warping",
      warper->Update() );

   FloatImageType3D::Pointer warpedImage = warper->GetOutput();


   {
      // compute final SSD 
      double finalSSD = 0.0;
      typedef itk::ImageRegionConstIterator<FloatImageType3D> ImageConstIterator;

      ImageConstIterator iterfix = ImageConstIterator(
         fixedImage, fixedImage->GetRequestedRegion() );

      ImageConstIterator itermovwarp = ImageConstIterator(
         warper->GetOutput(), fixedImage->GetRequestedRegion() );

      for (iterfix.GoToBegin(), itermovwarp.GoToBegin(); !iterfix.IsAtEnd(); ++iterfix, ++itermovwarp)
      {
         finalSSD += vnl_math_sqr( iterfix.Get() - itermovwarp.Get() );
      }

      const double finalMSE = finalSSD / static_cast<double>(
         fixedImage->GetRequestedRegion().GetNumberOfPixels() );

      std::cout << "MSE fixed image vs. warped moving image: " << finalMSE << std::endl;
   }



   // finally copy the warped output image back into the moving image
   ITK_IMAGE_COPY_MACRO_TN( FloatImageType3D, warpedImage, 
                            InputImageType, this->m_movingImage );
   
  
   // if necessary write the displacement field
   if (this->m_displacementFieldFilename != std::string(""))
   {
      typedef itk::ImageFileWriter<DeformationFieldType> DefFieldWriter;
      DefFieldWriter::Pointer def_field_writer = DefFieldWriter::New();
      def_field_writer->SetFileName(this->m_displacementFieldFilename.c_str());
      def_field_writer->SetInput( defField );
      std::cout << "writing displacement field: " << 
         this->m_displacementFieldFilename.c_str() << std::endl;
      def_field_writer->Update();
   }
   
   return CTNonlinearRegistration::OK;
}





#endif /*CTDIFFEOMORPHICDEMONSREGISTRATION_TXX_*/