/*
 * Copyright (c) ICG. All rights reserved.
 * See copyright.txt for more information.
 *
 * 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     : MIPItkProjects
 * Module      : Evaluation
 * Class       : $RCSfile $
 * Language    : C++
 * Description : 
 *
 * Author     : Martin Urschler
 * EMail      : urschler@icg.tu-graz.ac.at
 * Date       : $Date: $
 * Version    : $Revision:  $
 * Full Id    : $Id:  $
 *
 */


#include "itkExtractImageFilter.h"
#include "itkImageFileWriter.h"
#include "itkShiftScaleImageFilter.h"

#include "CTNonlinearRegistrationInterface.h"

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



#include "mlconfig.hpp"



int main(int argc, char**argv)
{  
   if (argc < 1 || argc > 2)
   {
      std::cerr << "Usage: " << argv[0] << "  config.txt" << std::endl;
      std::cerr << "Default config filename is 'config.txt', if it doesn't " <<
         " exist yet, please create it by copying from config_template.txt " <<
         " and modifying its entries!" << std::endl;
      return EXIT_FAILURE;
   }

   
   
   std::string configFilename = "config.txt";
   if (argc == 2)
      configFilename = std::string( argv[1] );

   // read and initialize config datastructure
   MLConfig::setupConfigFile(configFilename.c_str());   

   const std::string filenameFixed = 
      MLConfig::stringval( "FIXED_INPUT_FILENAME" );
   const std::string filenameMoving = 
      MLConfig::stringval( "MOVING_INPUT_FILENAME" );

   const Int16ImageType::PixelType resampleDefaultValue = 
      MLConfig::intval( "RESAMPLE_DEFAULT_VALUE" );
      
   //const bool preSmoothInputImages =
   //   MLConfig::boolval( "PRE_SMOOTH_INPUT_IMAGES" );

   const std::string mode = MLConfig::stringval( "REGISTRATION_MODE" );
   const bool writeDisplacementField = 
      MLConfig::boolval( "WRITE_DISPLACEMENT_FIELD" );

   std::string filenameOutputImage("");
   if (MLConfig::containsKey( "OUTPUT_IMAGE_FILENAME" ))
   {
      filenameOutputImage = MLConfig::stringval( "OUTPUT_IMAGE_FILENAME" ); 
   }
   
   std::string filenameOutputDefField("");
   if (MLConfig::containsKey( "OUTPUT_DISPLACEMENT_FILENAME" ))
   {
      filenameOutputDefField = 
         MLConfig::stringval( "OUTPUT_DISPLACEMENT_FILENAME" ); 
   }

   const std::string displacementExtensionString(".mha");

   Int16ImageType::Pointer fixedImage =
      VolumeIOWrapper<Int16ImageType>::readITKVolume( filenameFixed, "" );
   Int16ImageType::Pointer movingImage =
      VolumeIOWrapper<Int16ImageType>::readITKVolume( filenameMoving, "" );
   
   
   const Int16ImageType::SizeType fixedSize = 
      fixedImage->GetLargestPossibleRegion().GetSize();
   const Int16ImageType::SizeType movingSize = 
      movingImage->GetLargestPossibleRegion().GetSize();
    
   std::cout << "fixedSize: " << fixedSize << std::endl;
   std::cout << "movingSize: " << movingSize << std::endl;
   
   const Int16ImageType::SpacingType fixedSpacing  = fixedImage->GetSpacing();
   const Int16ImageType::SpacingType movingSpacing = movingImage->GetSpacing();
   
   std::cout << "fixedSpacing: " << fixedSpacing << std::endl;
   std::cout << "movingSpacing: " << movingSpacing << std::endl;
   
   
   //if (fixedSize != movingSize)
   //{
   //   std::cerr << "fixed and moving image size mismatch" << std::endl;
   //   return EXIT_FAILURE;
   //}
   
   
   signed short *baseptr = 
      new signed short[ (fixedSize[0]*fixedSize[1]*fixedSize[2]) ];
   signed short *floatingptr = 
      new signed short[ (movingSize[0]*movingSize[1]*movingSize[2]) ];
   
   
   const signed short fixedMinimum = 
      SmallUtilityMethods<Int16ImageType>::getMinimumVoxelValue(fixedImage);
   const signed short movingMinimum = 
      SmallUtilityMethods<Int16ImageType>::getMinimumVoxelValue(movingImage);

   std::cout << "fixed minimum: " << fixedMinimum << std::endl;
   std::cout << "moving minimum: " << movingMinimum << std::endl;
   
   {   
      typedef itk::ImageRegionConstIterator<Int16ImageType> ConstImageIterator;
      ConstImageIterator fixedIter( fixedImage, 
         fixedImage->GetLargestPossibleRegion() );

      unsigned long count = 0L;

      fixedIter.GoToBegin();
      while (!fixedIter.IsAtEnd())
      {
         baseptr[count] = fixedIter.Value();

         ++count;
         ++fixedIter;
      }
   }
   {   
      typedef itk::ImageRegionConstIterator<Int16ImageType> ConstImageIterator;
      ConstImageIterator movingIter( movingImage, 
         movingImage->GetLargestPossibleRegion() );

      unsigned long count = 0L;

      movingIter.GoToBegin();
      while (!movingIter.IsAtEnd())
      {
         floatingptr[count] = movingIter.Value();

         ++count;
         ++movingIter;
      }
   }
   
   // we don't need this ones anymore so we release their memory
   fixedImage = 0;
   
   std::cout << mode.c_str() << std::endl;
   
   itk::TimeProbesCollectorBase collector;
   
   if (mode == std::string("demons"))
   {
      std::string outputFilenameExtension = std::string("_unwarped_") + mode;
      const std::string timeCollectorString = mode + " registration";
      
      if (filenameOutputImage == std::string(""))
      {
         filenameOutputImage = filenameMoving;
         filenameOutputImage.replace( filenameOutputImage.find(".hdr"), 4,
            outputFilenameExtension + std::string(".hdr") );
      }
      
      collector.Start( timeCollectorString.c_str() );
      {  
         CTNonlinearRegistrationInterface<signed short> registration;
         registration.setResampleDefaultValue( resampleDefaultValue );
         
         const int nrLevels = 
            MLConfig::intval( "DEMONS_NR_LEVELS" );
         const unsigned int iter1 = 
            MLConfig::intval( "DEMONS_NR_ITERATIONS1" );
         const unsigned int iter2 = 
            MLConfig::intval( "DEMONS_NR_ITERATIONS2" );
         
         if (writeDisplacementField)
         {
            if (filenameOutputDefField != "")
            {
               registration.setDisplacementFieldFilename( 
                  filenameOutputDefField );
            }
            else
            {
               std::string displacementFieldFilename = filenameMoving;
               std::string dispFieldExt = outputFilenameExtension + 
                  displacementExtensionString;
               displacementFieldFilename.replace( 
                  displacementFieldFilename.find(".hdr"), 4, 
                  dispFieldExt.c_str() );
               registration.setDisplacementFieldFilename( 
                  displacementFieldFilename );
            }
         }
               
         registration.setImageFilenames( filenameFixed, filenameMoving );
         registration.setImageOutputFilename( filenameOutputImage );
         
         CTNonlinearRegistration::ReturnType return_val = 
            registration.standardDemonsMatching(
               baseptr, fixedSize[2], fixedSize[1], fixedSize[0], 
               fixedSpacing[2], fixedSpacing[0], fixedSpacing[1],
               floatingptr, movingSize[2], movingSize[1], movingSize[0], 
               movingSpacing[2], movingSpacing[0], movingSpacing[1],
               nrLevels, iter1, iter2 );
         
         std::cout << "return value: " << return_val << std::endl;
      }
      collector.Stop( timeCollectorString.c_str() );
      collector.Report();
   }
   else if (mode == std::string("levelset_motion"))
   {
      std::string outputFilenameExtension = std::string("_unwarped_") + mode;
      const std::string timeCollectorString = mode + " registration";
      
      if (filenameOutputImage == std::string(""))
      {
         filenameOutputImage = filenameMoving;
         filenameOutputImage.replace( filenameOutputImage.find(".hdr"), 4,
            outputFilenameExtension + std::string(".hdr") );
      }
      
      collector.Start( timeCollectorString.c_str() );
      {  
         CTNonlinearRegistrationInterface<signed short> registration;
         registration.setResampleDefaultValue( resampleDefaultValue );
         
         const int nrLevels = 
            MLConfig::intval( "LEVELSET_MOTION_NR_LEVELS" );
         const unsigned int iter1 = 
            MLConfig::intval( "LEVELSET_MOTION_NR_ITERATIONS1" );
         const unsigned int iter2 = 
            MLConfig::intval( "LEVELSET_MOTION_NR_ITERATIONS2" );
            
         if (writeDisplacementField)
         {
            if (filenameOutputDefField != "")
            {
               registration.setDisplacementFieldFilename( 
                  filenameOutputDefField );
            }
            else
            {
               std::string displacementFieldFilename = filenameMoving;
               std::string dispFieldExt = outputFilenameExtension + 
                  displacementExtensionString;
               displacementFieldFilename.replace( 
                  displacementFieldFilename.find(".hdr"), 4, 
                  dispFieldExt.c_str() );
               registration.setDisplacementFieldFilename( 
                  displacementFieldFilename );
            }
         }
               
         registration.setImageFilenames( filenameFixed, filenameMoving );
         registration.setImageOutputFilename( filenameOutputImage );
         
         CTNonlinearRegistration::ReturnType return_val = 
            registration.levelSetMotionMatching(
               baseptr, fixedSize[2], fixedSize[1], fixedSize[0], 
               fixedSpacing[2], fixedSpacing[0], fixedSpacing[1],
               floatingptr, movingSize[2], movingSize[1], movingSize[0], 
               movingSpacing[2], movingSpacing[0], movingSpacing[1],
               nrLevels, iter1, iter2 );
         
         std::cout << "return value: " << return_val << std::endl;
      }
      collector.Stop( timeCollectorString.c_str() );
      collector.Report();
   }

   else if (mode == std::string("symmetric_demons"))
   {
      std::string outputFilenameExtension = std::string("_unwarped_") + mode;
      const std::string timeCollectorString = mode + " registration";
      
      if (filenameOutputImage == std::string(""))
      {
         filenameOutputImage = filenameMoving;
         filenameOutputImage.replace( filenameOutputImage.find(".hdr"), 4,
            outputFilenameExtension + std::string(".hdr") );
      }
      
      collector.Start( timeCollectorString.c_str() );
      {  
         CTNonlinearRegistrationInterface<signed short> registration;
         registration.setResampleDefaultValue( resampleDefaultValue );
         
         const int nrLevels = 
            MLConfig::intval( "SYMMETRIC_DEMONS_NR_LEVELS" );
         const unsigned int iter1 = 
            MLConfig::intval( "SYMMETRIC_DEMONS_NR_ITERATIONS1" );
         const unsigned int iter2 = 
            MLConfig::intval( "SYMMETRIC_DEMONS_NR_ITERATIONS2" );
         
         if (writeDisplacementField)
         {
            if (filenameOutputDefField != "")
            {
               registration.setDisplacementFieldFilename( 
                  filenameOutputDefField );
            }
            else
            {
               std::string displacementFieldFilename = filenameMoving;
               std::string dispFieldExt = outputFilenameExtension + 
                  displacementExtensionString;
               displacementFieldFilename.replace( 
                  displacementFieldFilename.find(".hdr"), 4, 
                  dispFieldExt.c_str() );
               registration.setDisplacementFieldFilename( 
                  displacementFieldFilename );
            }
         }
               
         registration.setImageFilenames( filenameFixed, filenameMoving );
         registration.setImageOutputFilename( filenameOutputImage );
         
         CTNonlinearRegistration::ReturnType return_val = 
            registration.symmetricDemonsMatching(
               baseptr, fixedSize[2], fixedSize[1], fixedSize[0], 
               fixedSpacing[2], fixedSpacing[0], fixedSpacing[1],
               floatingptr, movingSize[2], movingSize[1], movingSize[0], 
               movingSpacing[2], movingSpacing[0], movingSpacing[1],
               nrLevels, iter1, iter2 );
         
         std::cout << "return value: " << return_val << std::endl;
      }
      collector.Stop( timeCollectorString.c_str() );
      collector.Report();
   }

   else if (mode == std::string("curvature"))
   {
      std::string outputFilenameExtension = std::string("_unwarped_") + mode;
      const std::string timeCollectorString = mode + " registration";
      
      if (filenameOutputImage == std::string(""))
      {
         filenameOutputImage = filenameMoving;
         filenameOutputImage.replace( filenameOutputImage.find(".hdr"), 4,
            outputFilenameExtension + std::string(".hdr") );
      }
      
      collector.Start( timeCollectorString.c_str() );
      {  
         CTNonlinearRegistrationInterface<signed short> registration;
         registration.setResampleDefaultValue( resampleDefaultValue );
         
         const int nrLevels = 
            MLConfig::intval( "CURVATURE_NR_LEVELS" );
         const unsigned int iter1 = 
            MLConfig::intval( "CURVATURE_NR_ITERATIONS1" );
         const unsigned int iter2 = 
            MLConfig::intval( "CURVATURE_NR_ITERATIONS2" );
         
         if (writeDisplacementField)
         {
            if (filenameOutputDefField != "")
            {
               registration.setDisplacementFieldFilename( 
                  filenameOutputDefField );
            }
            else
            {
               std::string displacementFieldFilename = filenameMoving;
               std::string dispFieldExt = outputFilenameExtension + 
                  displacementExtensionString;
               displacementFieldFilename.replace( 
                  displacementFieldFilename.find(".hdr"), 4, 
                  dispFieldExt.c_str() );
               registration.setDisplacementFieldFilename( 
                  displacementFieldFilename );
            }
         }
               
         registration.setImageFilenames( filenameFixed, filenameMoving );
         registration.setImageOutputFilename( filenameOutputImage );
         
         CTNonlinearRegistration::ReturnType return_val = 
            registration.curvatureMatching(
               baseptr, fixedSize[2], fixedSize[1], fixedSize[0], 
               fixedSpacing[2], fixedSpacing[0], fixedSpacing[1],
               floatingptr, movingSize[2], movingSize[1], movingSize[0], 
               movingSpacing[2], movingSpacing[0], movingSpacing[1],
               nrLevels, iter1, iter2 );
         
         std::cout << "return value: " << return_val << std::endl;
      }
      collector.Stop( timeCollectorString.c_str() );
      collector.Report();
   }


   else if (mode == std::string("fast_block_matching"))
   {
      std::string outputFilenameExtension = std::string("_unwarped_") + mode;
      const std::string timeCollectorString = mode + " registration";
      
      if (filenameOutputImage == std::string(""))
      {
         filenameOutputImage = filenameMoving;
         filenameOutputImage.replace( filenameOutputImage.find(".hdr"), 4,
            outputFilenameExtension + std::string(".hdr") );
      }
      
      collector.Start( timeCollectorString.c_str() );
      {  
         CTNonlinearRegistrationInterface<signed short> registration;
         registration.setResampleDefaultValue( resampleDefaultValue );
         
         const float regularization = 
            MLConfig::doubleval( "FAST_BLOCK_MATCHING_REGULARIZATION" );
         const float similarity = 
            MLConfig::doubleval( "FAST_BLOCK_MATCHING_SIMILARITY" );
         const unsigned int levelsNotToCompute = 
            MLConfig::intval( "FAST_BLOCK_MATCHING_LEVELS_NOT_TO_COMPUTE" );
         
         if (writeDisplacementField)
         {
            if (filenameOutputDefField != "")
            {
               registration.setDisplacementFieldFilename( 
                  filenameOutputDefField );
            }
            else
            {
               std::string displacementFieldFilename = filenameMoving;
               std::string dispFieldExt = outputFilenameExtension + 
                  displacementExtensionString;
               displacementFieldFilename.replace( 
                  displacementFieldFilename.find(".hdr"), 4, 
                  dispFieldExt.c_str() );
               registration.setDisplacementFieldFilename( 
                  displacementFieldFilename );
            }
         }
         
         registration.setImageFilenames( filenameFixed, filenameMoving );
         registration.setImageOutputFilename( filenameOutputImage );
         
         CTNonlinearRegistration::ReturnType return_val = 
            registration.fastBlockMatching(
               baseptr, fixedSize[2], fixedSize[1], fixedSize[0], 
               fixedSpacing[2], fixedSpacing[0], fixedSpacing[1],
               floatingptr, movingSize[2], movingSize[1], movingSize[0], 
               movingSpacing[2], movingSpacing[0], movingSpacing[1],
               regularization, similarity, levelsNotToCompute );
         
         std::cout << "return value: " << return_val << std::endl;
      }
      collector.Stop( timeCollectorString.c_str() );
      collector.Report();
   }
   else if (mode == std::string("diffeomorphic_demons"))
   {
      std::string outputFilenameExtension = std::string("_unwarped_") + mode;
      const std::string timeCollectorString = mode + " registration";
      
      if (filenameOutputImage == std::string(""))
      {
         filenameOutputImage = filenameMoving;
         filenameOutputImage.replace( filenameOutputImage.find(".hdr"), 4,
            outputFilenameExtension + std::string(".hdr") );
      }
      
      collector.Start( timeCollectorString.c_str() );
      {  
         CTNonlinearRegistrationInterface<signed short> registration;
         registration.setResampleDefaultValue( resampleDefaultValue );
         
         const unsigned int nrLevels = 
            MLConfig::intval( "DIFFEOMORPHIC_DEMONS_NR_LEVELS" );
         const unsigned int iters1 = 
            MLConfig::intval( "DIFFEOMORPHIC_DEMONS_NR_ITERATIONS1" );
         const unsigned int iters2 = 
            MLConfig::intval( "DIFFEOMORPHIC_DEMONS_NR_ITERATIONS2" );
         
         const float sigmaDeformationField =
            MLConfig::doubleval( "DIFFEOMORPHIC_DEMONS_SIGMA_DEFORMATION_FIELD" );
         const float sigmaUpdateField =
            MLConfig::doubleval( "DIFFEOMORPHIC_DEMONS_SIGMA_UPDATE_FIELD" );
         const float maxLengthUpdateField =
            MLConfig::doubleval( "DIFFEOMORPHIC_DEMONS_MAX_LENGTH_UPDATE_FIELD" );
         const unsigned int demonsForceGradientType =
            MLConfig::intval( "DIFFEOMORPHIC_DEMONS_FORCE_GRADIENT_TYPE" );
         
         if (writeDisplacementField)
         {
            if (filenameOutputDefField != "")
            {
               registration.setDisplacementFieldFilename( 
                  filenameOutputDefField );
            }
            else
            {
               std::string displacementFieldFilename = filenameMoving;
               std::string dispFieldExt = outputFilenameExtension + 
                  displacementExtensionString;
               displacementFieldFilename.replace( 
                  displacementFieldFilename.find(".hdr"), 4, 
                  dispFieldExt.c_str() );
               registration.setDisplacementFieldFilename( 
                  displacementFieldFilename );
            }
         }
         
         registration.setImageFilenames( filenameFixed, filenameMoving );
         registration.setImageOutputFilename( filenameOutputImage );
         
         CTNonlinearRegistration::ReturnType return_val = 
            registration.diffeomorphicDemonsMatching(
               baseptr, fixedSize[2], fixedSize[1], fixedSize[0], 
               fixedSpacing[2], fixedSpacing[0], fixedSpacing[1],
               floatingptr, movingSize[2], movingSize[1], movingSize[0], 
               movingSpacing[2], movingSpacing[0], movingSpacing[1],
               nrLevels, iters1, iters2, sigmaDeformationField,
               sigmaUpdateField, maxLengthUpdateField, 
               demonsForceGradientType );
         
         std::cout << "return value: " << return_val << std::endl;
      }
      collector.Stop( timeCollectorString.c_str() );
      collector.Report();
   }
   
   else
   {
      std::cerr << "Unknown registration mode: " << mode.c_str() << std::endl;
      delete [] baseptr;
      delete [] floatingptr;

      return EXIT_FAILURE;
   }
   
   
   {
      typedef itk::ImageRegionIterator<Int16ImageType> ImageIterator;
      ImageIterator movingIter( movingImage, 
         movingImage->GetLargestPossibleRegion() );

      unsigned long count = 0L;

      movingIter.GoToBegin();
      while (!movingIter.IsAtEnd())
      {
         movingIter.Set( floatingptr[count] );

         ++count;
         ++movingIter;
      }
   }
   
   VolumeIOWrapper<Int16ImageType>::writeITKVolume16Bit( movingImage, 
      filenameOutputImage, "" );
   
   delete [] baseptr;
   delete [] floatingptr;

   return EXIT_SUCCESS;
}