/*
* 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: CTNonlinearRegistrationBase.txx,v $
* Language    : C++
* Description : 
*
* Author     : Martin Urschler
* EMail      : urschler@icg.tu-graz.ac.at
* Date       : $Date: 2007-07-02 15:15:04 $
* Version    : $Revision: 1.22 $
* Full Id    : $Id: CTNonlinearRegistrationBase.txx,v 1.22 2007-07-02 15:15:04 urschler Exp $
*
*/


#ifndef _CTNONLINEARREGISTRATIONBASE_TXX_
#define _CTNONLINEARREGISTRATIONBASE_TXX_


#include "RawDataBufferImporter.h"
#include "VolumeIOWrapper.h"

#include "itkWarpImageFilter.h"
#include "itkImageFileWriter.h"


template<typename TPixelDataType>
CTNonlinearRegistrationBase<TPixelDataType>::
CTNonlinearRegistrationBase() 
   : m_fixedImage(0), m_movingImage(0), 
     m_displacementFieldFilename(""),
     m_imageOutputFilename(""),
     m_resampleDefaultValue(0)
{
}

template<typename TPixelDataType>
CTNonlinearRegistrationBase<TPixelDataType>::
~CTNonlinearRegistrationBase() 
{
   m_fixedImage = 0;
   m_movingImage = 0;
}

template<typename TPixelDataType>
void
CTNonlinearRegistrationBase<TPixelDataType>::
setDisplacementFieldFilename( 
   const std::string& displacementFieldFilename )
{
   m_displacementFieldFilename = displacementFieldFilename;
}

template<typename TPixelDataType>
void 
CTNonlinearRegistrationBase<TPixelDataType>::
setImageFilenames( const std::string& fixedFilename,
   const std::string& movingFilename )
{
   m_fixedFilename = fixedFilename;
   m_movingFilename = movingFilename;
}

template<typename TPixelDataType>
void
CTNonlinearRegistrationBase<TPixelDataType>::
setImageOutputFilename( const std::string& outputFilename )
{
   m_imageOutputFilename = (outputFilename == "") ? 
      m_movingFilename : outputFilename;
}

template<typename TPixelDataType>
void 
CTNonlinearRegistrationBase<TPixelDataType>::
setResampleDefaultValue( const int resampleDefaultValue )
{
   m_resampleDefaultValue = resampleDefaultValue;
}

template<typename TPixelDataType>
CTNonlinearRegistration::ReturnType
CTNonlinearRegistrationBase<TPixelDataType>::
setFixedImage(TPixelDataType * ptr, 
   const int nrSlices, const int nrRows, const int nrCols,
   const float sliceDistance, const float inplaneVoxelSizeX, 
   const float inplaneVoxelSizeY 
   )
{
   // now try to import the actual image data 
   EXCEPTION_CHECKED_ITK_FILTER_UPDATE( "Importing fixed image!",
      m_fixedImage = RawDataBufferImporter<TPixelDataType>::import( 
         ptr,
         nrSlices, nrRows, nrCols, 
         sliceDistance, inplaneVoxelSizeX, inplaneVoxelSizeY )
   );
   
   m_fixedImageSize = m_fixedImage->GetLargestPossibleRegion().GetSize();
   m_fixedImageSpacing = m_fixedImage->GetSpacing();
   
   // perform a gauss smoothing with sigma = 1 voxel
//   m_fixedImage = NonlinearRegistrationUtilities::
//      efficientDiscreteGaussSmoothVolume<InputImageType>( m_fixedImage,
//         m_fixedImageSpacing[0], 
//         m_fixedImageSpacing[1], 
//         m_fixedImageSpacing[2] );
   
   return CTNonlinearRegistration::OK;
}
      
template<typename TPixelDataType>
CTNonlinearRegistration::ReturnType
CTNonlinearRegistrationBase<TPixelDataType>::
setMovingImage(TPixelDataType * ptr, 
   const int nrSlices, const int nrRows, const int nrCols,
   const float sliceDistance, const float inplaneVoxelSizeX, 
   const float inplaneVoxelSizeY 
   )
{
   // now try to import the actual image data 
   EXCEPTION_CHECKED_ITK_FILTER_UPDATE( "Importing moving image!",
      m_movingImage = RawDataBufferImporter<TPixelDataType>::import( 
         ptr,
         nrSlices, nrRows, nrCols, 
         sliceDistance, inplaneVoxelSizeX, inplaneVoxelSizeY )
   );
   
   m_movingImageSize = m_movingImage->GetLargestPossibleRegion().GetSize();
   m_movingImageSpacing = m_movingImage->GetSpacing();
   
   // perform a gauss smoothing with sigma = 1 voxel
//   m_movingImage = NonlinearRegistrationUtilities::
//      efficientDiscreteGaussSmoothVolume<InputImageType>( m_movingImage,
//         m_movingImageSpacing[0], 
//         m_movingImageSpacing[1], 
//         m_movingImageSpacing[2] );
         
   return CTNonlinearRegistration::OK;
}
 
template<typename TPixelDataType>
CTNonlinearRegistration::ReturnType
CTNonlinearRegistrationBase<TPixelDataType>::
calculateFixedAndMovingPyramid( 
   PyramidType& fixedPyramid, PyramidType& movingPyramid )
{
   // first of all we look at the image sizes to be able to argue about
   // the downsample strategy ( and the associated upsample strategy )
   
   typename InputImageType::SizeType fixedFullSize = 
      m_fixedImage->GetLargestPossibleRegion().GetSize();
   typename InputImageType::SizeType movingFullSize = 
      m_movingImage->GetLargestPossibleRegion().GetSize();
   
   // FIXXME: this is very restrictive - find a weaker condition by adapting
   // the algorithms
   if (fixedFullSize != movingFullSize)
   {
      return CTNonlinearRegistration::VOLUME_SIZES_MISMATCH_EXCEPTION;
   }
   
//   if( testPowerOfTwoImageSize(fixedFullSize)  == false ||
//       testPowerOfTwoImageSize(movingFullSize) == false )
//   {
//      return CTNonlinearRegistration::VOLUME_SIZES_NOT_POWER_OF_TWO_EXCEPTION;
//   }
   
   // the pyramid includes all scale levels, also the original one, therefore -1
   const unsigned int nrDivisions = fixedPyramid.size() - 1;
   if( testDivisibleByTwoImageSize(fixedFullSize,nrDivisions)  == false ||
       testDivisibleByTwoImageSize(movingFullSize,nrDivisions) == false )
   {
      return CTNonlinearRegistration::
         VOLUME_SIZES_NOT_DIVISIBLE_BY_TWO_ENOUGH_TIMES_EXCEPTION;
   }
   
   CTNonlinearRegistration::ReturnType pyramidReturnValue =
      calculateImagePyramid( m_fixedImage, fixedPyramid, 
                             m_fixedDownsampleModeSequence );
   
   if (pyramidReturnValue != CTNonlinearRegistration::OK)
      return pyramidReturnValue;
   
   pyramidReturnValue = calculateImagePyramid( m_movingImage, movingPyramid,
                                               m_movingDownsampleModeSequence );
   
   if (pyramidReturnValue != CTNonlinearRegistration::OK)
      return pyramidReturnValue;
   
#if 0
   const unsigned int maxNrShrinkLevels = fixedPyramid.size();
   
   for ( unsigned int pyramidLevel=1; pyramidLevel<maxNrShrinkLevels; 
         ++pyramidLevel)
   {
      std::string outputExtensionDiff;
      std::string outputExtensionDiff2;
      if (pyramidLevel==1)
      {
         outputExtensionDiff = std::string("_shrinked1");
         outputExtensionDiff2 = std::string("_shrinked_upsampled1");
      }
      else if (pyramidLevel==2)
      {
         outputExtensionDiff = std::string("_shrinked2");
         outputExtensionDiff2 = std::string("_shrinked_upsampled2");
      }
      else if (pyramidLevel==3)
      {     
         outputExtensionDiff = std::string("_shrinked3");
         outputExtensionDiff2 = std::string("_shrinked_upsampled3");
      }
      else if (pyramidLevel==4)
      {
         outputExtensionDiff = std::string("_shrinked4");
         outputExtensionDiff2 = std::string("_shrinked_upsampled4");
      }
      else if (pyramidLevel==5)
      {
         outputExtensionDiff = std::string("_shrinked5");
         outputExtensionDiff2 = std::string("_shrinked_upsampled5");
      }
      else if (pyramidLevel==6)
      {
         outputExtensionDiff = std::string("_shrinked6");
         outputExtensionDiff2 = std::string("_shrinked_upsampled6");
      }
      
      if (pyramidLevel<=6)
      {
         VolumeIOWrapper<InputImageType>::writeITKVolume16Bit( 
            fixedPyramid[pyramidLevel], m_fixedFilename, 
            outputExtensionDiff );
         VolumeIOWrapper<InputImageType>::writeITKVolume16Bit( 
            movingPyramid[pyramidLevel], m_movingFilename, 
            outputExtensionDiff );
         
         // upsample and store the original resolution
         InputImageType::Pointer tmpImage = fixedPyramid[pyramidLevel];
         for (unsigned int i=0;i<pyramidLevel;++i)
         {
            tmpImage = DownAndUpsampleUtilities::
               doubleImageByLinearInterpolation<InputImageType>(tmpImage);
         }
         VolumeIOWrapper<InputImageType>::writeITKVolume16Bit( 
            tmpImage, m_fixedFilename, outputExtensionDiff2 );
         
         tmpImage = movingPyramid[pyramidLevel];
         for (unsigned int i=0;i<pyramidLevel;++i)
         {
            tmpImage = DownAndUpsampleUtilities::
               doubleImageByLinearInterpolation<InputImageType>(tmpImage);
         }
         VolumeIOWrapper<InputImageType>::writeITKVolume16Bit( 
            tmpImage, m_movingFilename, outputExtensionDiff2 );
      }
   }
#endif
   
   return CTNonlinearRegistration::OK;
}

template<typename TPixelDataType>
CTNonlinearRegistration::ReturnType
CTNonlinearRegistrationBase<TPixelDataType>::
calculateImagePyramid( 
   const typename InputImageType::Pointer& image,
   PyramidType& pyramid,
   DownAndUpsampleUtilities::PyramidDownsampleModeSequence& 
      downsampleModeSequence )
{
   // initialize the pyramid
   pyramid[0]  = image;
   
   const unsigned int maxNrShrinkLevels = pyramid.size();
   
   for ( unsigned int pyramidLevel=1; pyramidLevel<maxNrShrinkLevels; 
         ++pyramidLevel)
   {
      // calculate the downsampled pyramid levels
      std::cout << "pyramid level " << pyramidLevel << " ... " << std::endl;
      
      typename InputImageType::Pointer shrinked = 0;
      
      DownAndUpsampleUtilities::PyramidDownsampleModeNTuple
         downsampleModeNTuple;
      
      //EXCEPTION_CHECKED_ITK_FILTER_UPDATE("shrinking image",
      shrinked = 
         DownAndUpsampleUtilities::shrinkImageToHalfByBinomialSmoothing<
            InputImageType,InputImageType>( 
               pyramid[pyramidLevel-1], downsampleModeNTuple );

      pyramid[pyramidLevel] = shrinked;
      
      downsampleModeSequence.push_back( downsampleModeNTuple );
      
      std::cout << "   ... finished " << std::endl;
   }
   
   return CTNonlinearRegistration::OK;
}

template<typename TPixelDataType>
CTNonlinearRegistration::ReturnType
CTNonlinearRegistrationBase<TPixelDataType>::
warpMovingImage( const DeformationFieldType::Pointer deformationField )
{
   typedef itk::WarpImageFilter< InputImageType, InputImageType, 
      DeformationFieldType > WarperType;
   typedef itk::LinearInterpolateImageFunction< InputImageType, double >
      InterpolatorType;

   typename WarperType::Pointer warper = WarperType::New();
   typename InterpolatorType::Pointer interpolator = InterpolatorType::New();

   warper->SetInput( m_movingImage );
   warper->SetInterpolator( interpolator );
   warper->SetOutputSpacing( m_fixedImage->GetSpacing() );
   warper->SetOutputOrigin(  m_fixedImage->GetOrigin() );
   warper->SetEdgePaddingValue( m_resampleDefaultValue );
   
   warper->SetDeformationField( deformationField );
   
   EXCEPTION_CHECKED_ITK_FILTER_UPDATE("perform warping",
      warper->Update() );
      
   typename InputImageType::Pointer outputImage = warper->GetOutput();
   
   // finally copy the warped output image back into the moving image
   ITK_IMAGE_COPY_MACRO_TN( InputImageType, outputImage, 
                            InputImageType, m_movingImage );

   // if necessary write the displacement field
   if (m_displacementFieldFilename != std::string(""))
   {
      typedef itk::ImageFileWriter<DeformationFieldType> DefFieldWriter;
      DefFieldWriter::Pointer def_field_writer = DefFieldWriter::New();
      def_field_writer->SetFileName(m_displacementFieldFilename.c_str());
      def_field_writer->SetInput( deformationField );
      std::cout << "writing displacement field: " << 
         m_displacementFieldFilename.c_str() << std::endl;
      def_field_writer->Update();   
   }

   return CTNonlinearRegistration::OK;
}

template<typename TPixelDataType>
bool
CTNonlinearRegistrationBase<TPixelDataType>::
testPowerOfTwoImageSize( const typename InputImageType::SizeType& size )
{
   bool testOK = true;
   
   const unsigned int dimensions = 
      itk::GetImageDimension<InputImageType>::ImageDimension;
      
   for (unsigned int dim=0; dim<dimensions; ++dim)
   {
      // even image size
      if (size[dim] % 2 == 0)
      {
         const unsigned int n = size[dim];
         if ( (n&(n-1)) != 0 )
            testOK = false;
      }
      else // odd image size
      {
         const unsigned int n = size[dim]-1;
         if ( (n&(n-1)) != 0 )
            testOK = false;
      }
   }
   
   return testOK;
}

template<typename TPixelDataType>
bool
CTNonlinearRegistrationBase<TPixelDataType>::
testDivisibleByTwoImageSize( 
   const typename InputImageType::SizeType& size,
   const unsigned int nrTimesDivisible )
{
   bool testOK = true;
   
   // trivial check
   if (nrTimesDivisible == 0)
      return testOK;
   
   
   const unsigned int dimensions = 
      itk::GetImageDimension<InputImageType>::ImageDimension;
   
   const unsigned int bitMask = ((1 << nrTimesDivisible) - 1);
   
   //std::cout << "nrTimesDivisible = " << nrTimesDivisible << " - bitMask: " <<
   //   bitMask << std::endl; 
   
   for (unsigned int dim=0; dim<dimensions; ++dim)
   {
      // even image size
      if (size[dim] % 2 == 0)
      {
         const unsigned int n = size[dim];
         if ((n & bitMask) != 0)
            testOK = false;
      }
      else // odd image size
      {
         const unsigned int n = size[dim]-1;
         if ((n & bitMask) != 0)
            testOK = false;
      }
   }
   
   return testOK;
}

#endif // _CTNONLINEARREGISTRATIONBASE_TXX_