/*
 * 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: ApplySimulatedBreathingTransformationWithIntensityVariation.cxx,v $
 * Language    : C++
 * Description : 
 *
 * Author     : Martin Urschler
 * EMail      : urschler@icg.tu-graz.ac.at
 * Date       : $Date: 2007-03-07 10:01:10 $
 * Version    : $Revision: 1.13 $
 * Full Id    : $Id: ApplySimulatedBreathingTransformationWithIntensityVariation.cxx,v 1.13 2007-03-07 10:01:10 urschler Exp $
 *
 */



// include files

#include "commonItkTypedefs.h"
#include "SmallUtilityMethods.h"
#include "VolumeIOWrapper.h"
#include "SimulatedBreathingTransformation.h"

#include "itkImageFileWriter.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkNearestNeighborInterpolateImageFunction.h"
#include "vnl/vnl_random.h"


/*
 * This method simulates breathing by a simple transformation.
 * Rib cage movement is simulated using a radial inwards directed force.
 * Diaphragm movement is simulated using a force directed in the negative z 
 * direction, which is stronger near the center and less strong in the outer 
 * regions of the x-y plane.
 * The distribution of the latter force is calculated according to a Gaussian.
 * 
 * The deformation is intended to be applied to an inhalation data set and will 
 * lead to a simulated exhalation data set.
 * 
 * Additionally an intensity transformation is performed in the lung region 
 * (determined by simple thresholding). The intensities are modified randomly.
*/

// Note: it is currently crucial, that the original volume and the output volume
//       have exactly the same size, spacing and origin! Otherwise the
//       deformation field is invalid. This is ensured, since we create the 
//       output image with the original image as a template!

int main( int argc, char *argv[] )
{
   if (argc != 6)
   {
      std::cerr << "Usage: " << std::endl;
      std::cerr << argv[0] << 
         "  original_volume  output_volume"
         "  translation_vertical_mm  translation_inplane_mm" << 
         "  resample_default_value" << std::endl;
      return EXIT_FAILURE;
   }
   
   vnl_random randomNumberGenerator;
   
   const std::string input_filename( argv[1] );
   const std::string output_filename( argv[2] );

   // read the images
   Int16ImageType::Pointer original_image =
      VolumeIOWrapper<Int16ImageType>::readITKVolume( input_filename, "" );
   Int16ImageType::SizeType orig_size = original_image->GetLargestPossibleRegion().GetSize();
   Int16ImageType::SpacingType spacing = original_image->GetSpacing();
   
   const double translation_vertical_mm = atof( argv[3] );
   const double translation_inplane_mm = atof( argv[4] );

   const Int16ImageType::PixelType resampleDefaultValue = atoi( argv[5] );
   
   // remove the background
   MaskImageType::Pointer segmentation_image = 0;
   SmallUtilityMethods<Int16ImageType>::removeBackground( 
      original_image, segmentation_image, -1, -1500, -250 );

   //VolumeIOWrapper<MaskImageType>::writeITKVolume16Bit( segmentation_image, 
   //   input_filename, "_background_removal"  );

   // the warped segmentation image
   MaskImageType::Pointer warped_segmentation_image = 
      SmallUtilityMethods<MaskImageType>::createNewImage( segmentation_image );
   warped_segmentation_image->FillBuffer( 0 );
   
   
   // create a transformation object
   SimulatedBreathingTransformation synth_transf( segmentation_image, 
      translation_vertical_mm, translation_inplane_mm );

   Int16ImageType::Pointer output_image = 
      SmallUtilityMethods<Int16ImageType>::createNewImage( original_image );

   // the deformation field is defined according to the output image! 
   DeformationFieldType::Pointer deformation_field = 
      SmallUtilityMethods<Int16ImageType>::createNewZeroedDeformationField( 
         output_image );


   // this procedure is similar to the itkWarpImageFilter
   
   typedef itk::LinearInterpolateImageFunction<Int16ImageType, double>
      InterpolatorType;
   InterpolatorType::Pointer interpolator = InterpolatorType::New();
   interpolator->SetInputImage( original_image );

   typedef itk::NearestNeighborInterpolateImageFunction<MaskImageType, double> 
      MaskInterpolatorType;
   MaskInterpolatorType::Pointer mask_interpolator =
      MaskInterpolatorType::New();
   mask_interpolator->SetInputImage( segmentation_image );

   // iterator for the output image
   itk::ImageRegionIteratorWithIndex<Int16ImageType> out_it( output_image,
      output_image->GetLargestPossibleRegion() );
   
   // iterator for the output segmentation mask
   itk::ImageRegionIteratorWithIndex<MaskImageType> out_segmask_it( 
      warped_segmentation_image,
      warped_segmentation_image->GetLargestPossibleRegion() );
         
   // iterator for the deformation field
   itk::ImageRegionIterator<DeformationFieldType> field_it(
      deformation_field, deformation_field->GetLargestPossibleRegion() );

   Int16ImageType::IndexType out_index;
   Int16ImageType::PointType point;
   Int16ImageType::PointType transformed_point;
   VectorPixelType deformation;

   while ( !out_it.IsAtEnd() )
   {
      // get the output image index
      out_index = out_it.GetIndex();
      output_image->TransformIndexToPhysicalPoint( out_index, point );
      
      transformed_point = synth_transf.apply( point );
      
      if ( mask_interpolator->IsInsideBuffer( transformed_point ) )
      {
         out_segmask_it.Set( static_cast<MaskImageType::PixelType>( 
            mask_interpolator->Evaluate( transformed_point ) ) );
      }
      
      if ( interpolator->IsInsideBuffer( transformed_point ) )
      {
         out_it.Set( static_cast<Int16ImageType::PixelType>(
            interpolator->Evaluate( transformed_point ) ) );
         
         deformation = transformed_point - point;
         field_it.Set( deformation );
      }
      else
      {
         out_it.Set( resampleDefaultValue );
         field_it.Set( 
            VectorPixelConstants<VectorPixelType::ValueType, 3>::undefined() );
      }
      
      
      ++out_it;
      ++field_it;
   }


   std::string displacement_field_filename = output_filename;
   displacement_field_filename.replace(
      displacement_field_filename.find(".hdr"), 4, ".mha" ); 
   typedef itk::ImageFileWriter<DeformationFieldType> DeformationFieldWriter;
   DeformationFieldWriter::Pointer writer = DeformationFieldWriter::New();
   writer->SetFileName( displacement_field_filename.c_str() );
   writer->SetInput( deformation_field );
   std::cout << "Writing deformation field!" << std::endl;
   writer->Update();

   // HERE THE INTENSITY VARIATION COMES INTO PLAY
   
   Int16ImageType::PixelType output_min_value = 
      SmallUtilityMethods<Int16ImageType>::getMinimumVoxelValue(output_image);
   
   std::cout << "minimum of transformed image: " << output_min_value 
      << std::endl;
      
   out_it.GoToBegin();
   out_segmask_it.GoToBegin();
   while ( !out_it.IsAtEnd() )
   {
      if (out_segmask_it.Get() > 0)
      {
         //std::cout << "gotcha" << std::endl;
         // if intensity lies between smallest value and -800 HU
         const Int16ImageType::PixelType voxelIntensity = out_it.Get();
         if (voxelIntensity >= output_min_value && voxelIntensity <= -800)
         {
            int intensity_offset = 
               25 + ROUND_DBL<int>(5.0*randomNumberGenerator.normal());
            //std::cout << "intensity_offset: " << intensity_offset 
            //   << std::endl;
            out_it.Set( voxelIntensity + intensity_offset );
         }
      }
      
      ++out_it;
      ++out_segmask_it;
   }

   VolumeIOWrapper<Int16ImageType>::writeITKVolume16Bit( output_image, 
      output_filename, ""  );
   

   return EXIT_SUCCESS;
}




/*// the old version
   itk::ImageRegionIterator<DeformationFieldType> def_it( deformation_field, 
      deformation_field->GetLargestPossibleRegion() );
   itk::ImageRegionIteratorWithIndex<Int16ImageType> out_it( output_image, 
      output_image->GetLargestPossibleRegion() );
   itk::ImageRegionConstIteratorWithIndex<Int16ImageType> inp_it( 
      original_image, original_image->GetLargestPossibleRegion() );
   inp_it.GoToBegin();
   out_it.GoToBegin();
   def_it.GoToBegin();
   while ( !inp_it.IsAtEnd() )
   {
      const Int16ImageType::IndexType index = inp_it.GetIndex();

      PointSetType::PointType p;
      p[0] = index[0];
      p[1] = index[1];
      p[2] = index[2];

      PointSetType::PointType q = synth_transf.apply( p );

      MaskImageType::IndexType transformed_index;
      transformed_index[0] = ROUND_DBL<MaskImageType::IndexValueType>( q[0] );
      transformed_index[1] = ROUND_DBL<MaskImageType::IndexValueType>( q[1] );
      transformed_index[2] = ROUND_DBL<MaskImageType::IndexValueType>( q[2] );

      // check if we aren't out of bounds!
      if ( (transformed_index[0] >= 0) && (transformed_index[1] >= 0) && 
           (transformed_index[2] >= 0) && (transformed_index[0] < orig_size[0]) 
           && (transformed_index[1] < orig_size[1]) && 
           (transformed_index[2] < orig_size[2]) )
      {
         out_it.SetIndex( transformed_index );
         out_it.Set( inp_it.Get() );

         VectorPixelType deformation;
         deformation[0] = (q[0] - p[0]) * spacing[0];
         deformation[1] = (q[1] - p[1]) * spacing[1];
         deformation[2] = (q[2] - p[2]) * spacing[2];
         def_it.Set( deformation );
      }

      ++inp_it;
      ++def_it;
   }*/




/*// old version
   itk::ImageRegionIteratorWithIndex<MaskImageType> mask_out_it( 
      warped_segmentation_image, 
      warped_segmentation_image->GetLargestPossibleRegion() );
   itk::ImageRegionConstIteratorWithIndex<MaskImageType> mask_inp_it( 
      segmentation_image, segmentation_image->GetLargestPossibleRegion() );
   mask_inp_it.GoToBegin();
   mask_out_it.GoToBegin();
   while ( !mask_inp_it.IsAtEnd() )
   {
      const Int16ImageType::IndexType index = mask_inp_it.GetIndex();

      PointSetType::PointType p;
      p[0] = index[0];
      p[1] = index[1];
      p[2] = index[2];

      PointSetType::PointType q = synth_transf.apply( p );

      MaskImageType::IndexType transformed_index;
      transformed_index[0] = ROUND_DBL<MaskImageType::IndexValueType>( q[0] );
      transformed_index[1] = ROUND_DBL<MaskImageType::IndexValueType>( q[1] );
      transformed_index[2] = ROUND_DBL<MaskImageType::IndexValueType>( q[2] );

      // check if we aren't out of bounds!
      if ( (transformed_index[0] >= 0) && (transformed_index[1] >= 0) && 
           (transformed_index[2] >= 0) && (transformed_index[0] < orig_size[0]) 
           && (transformed_index[1] < orig_size[1]) && 
           (transformed_index[2] < orig_size[2]) )
      {
         mask_out_it.SetIndex( transformed_index );
         mask_out_it.Set( mask_inp_it.Get() );
      }

      ++mask_inp_it;
   }
   
   //VolumeIOWrapper<MaskImageType>::writeITKVolume8Bit( 
   //   warped_segmentation_image, 
   //   output_filename, "_background_removal", false  );
   
   Int16ImageType::PixelType output_min_value = 
      SmallUtilityMethods<Int16ImageType>::getMinimumVoxelValue(output_image);
   
   std::cout << "displaced image minimum: " << output_min_value << std::endl;
   itk::ImageRegionIteratorWithIndex<MaskImageType> warped_seg_it( 
      warped_segmentation_image, 
      warped_segmentation_image->GetLargestPossibleRegion() );
      
   out_it.GoToBegin();
   seg_it.GoToBegin();
   while ( !out_it.IsAtEnd() )
   {
      if (seg_it.Get() > 0)
      {
         //std::cout << "gotcha" << std::endl;
         // if intensity lies between smallest value and -800 HU
         Int16ImageType::PixelType voxelIntensity = out_it.Get();
         if (voxelIntensity >= output_min_value && voxelIntensity <= -800)
         {
            int intensity_offset = 
               25 + ROUND_DBL<int>(5.0*randomNumberGenerator.normal());
            //std::cout << "intensity_offset: " << intensity_offset << std::endl;
            out_it.Set( voxelIntensity + intensity_offset );
         }
      }
      
      ++out_it;
      ++seg_it;
   }

   VolumeIOWrapper<Int16ImageType>::writeITKVolume16Bit( output_image, 
      output_filename, ""  );*/