/*=========================================================================
  Program:   vv                     http://www.creatis.insa-lyon.fr/rio/vv

  Authors belong to:
  - University of LYON              http://www.universite-lyon.fr/
  - Léon Bérard cancer center       http://www.centreleonberard.fr
  - CREATIS CNRS laboratory         http://www.creatis.insa-lyon.fr

  This software is distributed WITHOUT ANY WARRANTY; without even
  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
  PURPOSE.  See the copyright notices for more information.

  It is distributed under dual licence

  - BSD        See included LICENSE.txt file
  - CeCILL-B   http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
===========================================================================**/
#ifndef clitkAffineTransformGenericFilter_txx
#define clitkAffineTransformGenericFilter_txx

/* =================================================
 * @file   clitkAffineTransformGenericFilter.txx
 * @author
 * @date
 *
 * @brief
 *
 ===================================================*/


namespace clitk
{

//-----------------------------------------------------------
// Constructor
//-----------------------------------------------------------
template<class args_info_type>
AffineTransformGenericFilter<args_info_type>::AffineTransformGenericFilter()
{
  m_Verbose=false;
  m_InputFileName="";
}


//-----------------------------------------------------------
// Update
//-----------------------------------------------------------
template<class args_info_type>
void AffineTransformGenericFilter<args_info_type>::Update()
{
  // Read the Dimension and PixelType
  int Dimension, Components;
  std::string PixelType;
  ReadImageDimensionAndPixelType(m_InputFileName, Dimension, PixelType, Components);


  // Call UpdateWithDim
  if(Dimension==2) UpdateWithDim<2>(PixelType, Components);
  else if(Dimension==3) UpdateWithDim<3>(PixelType, Components);
  else if (Dimension==4)UpdateWithDim<4>(PixelType, Components);
  else {
    std::cout<<"Error, Only for 2, 3 or 4  Dimensions!!!"<<std::endl ;
    return;
  }
}

//-------------------------------------------------------------------
// Update with the number of dimensions
//-------------------------------------------------------------------
template<class args_info_type>
template<unsigned int Dimension>
void
AffineTransformGenericFilter<args_info_type>::UpdateWithDim(std::string PixelType, int Components)
{
  if (m_Verbose) std::cout << "Image was detected to be "<<Dimension<<"D and "<<Components<<" component(s) of "<<  PixelType<<"..."<<std::endl;

  if (Components==1) {
    if(PixelType == "short") {
      if (m_Verbose) std::cout << "Launching filter in "<< Dimension <<"D and signed short..." << std::endl;
      UpdateWithDimAndPixelType<Dimension, signed short>();
    }
    //    else if(PixelType == "unsigned_short"){
    //       if (m_Verbose) std::cout  << "Launching filter in "<< Dimension <<"D and unsigned_short..." << std::endl;
    //       UpdateWithDimAndPixelType<Dimension, unsigned short>();
    //     }

    else if (PixelType == "unsigned_char") {
      if (m_Verbose) std::cout  << "Launching filter in "<< Dimension <<"D and unsigned_char..." << std::endl;
      UpdateWithDimAndPixelType<Dimension, unsigned char>();
    }

    //     else if (PixelType == "char"){
    //       if (m_Verbose) std::cout  << "Launching filter in "<< Dimension <<"D and signed_char..." << std::endl;
    //       UpdateWithDimAndPixelType<Dimension, signed char>();
    //     }
    else {
      if (m_Verbose) std::cout  << "Launching filter in "<< Dimension <<"D and float..." << std::endl;
      UpdateWithDimAndPixelType<Dimension, float>();
    }
  }

  else if (Components==3) {
    if (m_Verbose) std::cout  << "Launching transform in "<< Dimension <<"D and 3D float (DVF)" << std::endl;
    UpdateWithDimAndVectorType<Dimension, itk::Vector<float, Dimension> >();
  }

  else std::cerr<<"Number of components is "<<Components<<", not supported!"<<std::endl;

}


//-------------------------------------------------------------------
// Update with the number of dimensions and the pixeltype
//-------------------------------------------------------------------
template<class args_info_type>
template <unsigned int Dimension, class  PixelType>
void
AffineTransformGenericFilter<args_info_type>::UpdateWithDimAndPixelType()
{

  // ImageTypes
  typedef itk::Image<PixelType, Dimension> InputImageType;
  typedef itk::Image<PixelType, Dimension> OutputImageType;

  // Read the input
  typedef itk::ImageFileReader<InputImageType> InputReaderType;
  typename InputReaderType::Pointer reader = InputReaderType::New();
  reader->SetFileName( m_InputFileName);
  reader->Update();
  typename InputImageType::Pointer input= reader->GetOutput();

  //Filter
  typedef  itk::ResampleImageFilter< InputImageType,OutputImageType >  ResampleFilterType;
  typename ResampleFilterType::Pointer resampler = ResampleFilterType::New();

  // Matrix
  typename itk::Matrix<double, Dimension+1, Dimension+1> matrix;
  if (m_ArgsInfo.rotate_given || m_ArgsInfo.translate_given)
  {
    if (m_ArgsInfo.matrix_given)
    {
      std::cerr << "You must use either rotate/translate or matrix options" << std::cout;
      return;
    }
    itk::Array<double> transformParameters(2 * Dimension);
    transformParameters.Fill(0.0);
    if (m_ArgsInfo.rotate_given)
    {
      if (Dimension == 2)
        transformParameters[0] = m_ArgsInfo.rotate_arg[0];
      else
        for (unsigned int i = 0; i < 3; i++)
          transformParameters[i] = m_ArgsInfo.rotate_arg[i];
    }
    if (m_ArgsInfo.translate_given)
    {
      int pos = 3;
      if (Dimension == 2)
        pos = 1;
      for (unsigned int i = 0; i < Dimension && i < 3; i++)
        transformParameters[pos++] = m_ArgsInfo.translate_arg[i];
    }
    if (Dimension == 4)
    {
      matrix.SetIdentity();
      itk::Matrix<double, 4, 4> tmp = GetForwardAffineMatrix3D(transformParameters);
      for (unsigned int i = 0; i < 3; ++i)
        for (unsigned int j = 0; j < 3; ++j)
          matrix[i][j] = tmp[i][j];
      for (unsigned int i = 0; i < 3; ++i)
        matrix[i][4] = tmp[i][3];
    }
    else
      matrix = GetForwardAffineMatrix<Dimension>(transformParameters);
  }
  else
  {
    if (m_ArgsInfo.matrix_given)
    {
      matrix= clitk::ReadMatrix<Dimension>(m_ArgsInfo.matrix_arg);
      if (m_Verbose) std::cout << "Reading the matrix..." << std::endl;
    }
    else
      matrix.SetIdentity();
  }
  if (m_Verbose)
    std::cout << "Using the following matrix:" << std::endl
              << matrix << std::endl;
  typename itk::Matrix<double, Dimension, Dimension> rotationMatrix = clitk::GetRotationalPartMatrix(matrix);
  typename itk::Vector<double, Dimension> translationPart = clitk::GetTranslationPartMatrix(matrix);

  // Transform
  typedef itk::AffineTransform<double, Dimension> AffineTransformType;
  typename AffineTransformType::Pointer affineTransform=AffineTransformType::New();
  affineTransform->SetMatrix(rotationMatrix);
  affineTransform->SetTranslation(translationPart);

  // Interp
  typedef clitk::GenericInterpolator<args_info_type, InputImageType, double> GenericInterpolatorType;
  typename GenericInterpolatorType::Pointer genericInterpolator=GenericInterpolatorType::New();
  genericInterpolator->SetArgsInfo(m_ArgsInfo);

  // Properties
  if (m_ArgsInfo.like_given) {
    typename InputReaderType::Pointer likeReader=InputReaderType::New();
    likeReader->SetFileName(m_ArgsInfo.like_arg);
    likeReader->Update();
    resampler->SetOutputParametersFromImage(likeReader->GetOutput());
  } else if(m_ArgsInfo.transform_grid_flag) {
    typename itk::Matrix<double, Dimension+1, Dimension+1> invMatrix( matrix.GetInverse() );
    typename itk::Matrix<double, Dimension, Dimension> invRotMatrix( clitk::GetRotationalPartMatrix(invMatrix) );
    typename itk::Vector<double,Dimension> invTrans =  clitk::GetTranslationPartMatrix(invMatrix);

    // Spacing is influenced by affine transform matrix and input direction
    typename InputImageType::SpacingType outputSpacing;
    outputSpacing = invRotMatrix *
                    input->GetDirection() *
                    input->GetSpacing();

    // Origin is influenced by translation but not by input direction
    typename InputImageType::PointType outputOrigin;
    outputOrigin = invRotMatrix *
                   input->GetOrigin() +
                   invTrans;

    // Size is influenced by affine transform matrix and input direction
    // Size is converted to double, transformed and converted back to size type.
    vnl_vector<double> vnlOutputSize(Dimension);
    for(unsigned int i=0; i< Dimension; i++) {
      vnlOutputSize[i] = input->GetLargestPossibleRegion().GetSize()[i];
    }
    vnlOutputSize = invRotMatrix *
                    input->GetDirection().GetVnlMatrix() *
                    vnlOutputSize;
    typename OutputImageType::SizeType outputSize;
    for(unsigned int i=0; i< Dimension; i++) {
      // If the size is negative, we have a flip and we must modify
      // the origin and the spacing accordingly.
      if(vnlOutputSize[i]<0.) {
        vnlOutputSize[i] *= -1.;
        outputOrigin[i] = outputOrigin[i] + outputSpacing[i] * (vnlOutputSize[i]-1);
        outputSpacing[i] *= -1.;
      }
      outputSize[i] = lrint(vnlOutputSize[i]);
    }
    resampler->SetSize( outputSize );
    resampler->SetOutputSpacing( outputSpacing );
    resampler->SetOutputOrigin( outputOrigin );
  } else {
    //Size
    typename OutputImageType::SizeType outputSize;
    if (m_ArgsInfo.size_given) {
      for(unsigned int i=0; i< Dimension; i++)
        outputSize[i]=m_ArgsInfo.size_arg[i];
    } else outputSize=input->GetLargestPossibleRegion().GetSize();

    //Spacing
    typename OutputImageType::SpacingType outputSpacing;
    if (m_ArgsInfo.spacing_given) {
      for(unsigned int i=0; i< Dimension; i++)
        outputSpacing[i]=m_ArgsInfo.spacing_arg[i];
    } else outputSpacing=input->GetSpacing();

    //Origin
    typename OutputImageType::PointType outputOrigin;
    if (m_ArgsInfo.origin_given) {
      for(unsigned int i=0; i< Dimension; i++)
        outputOrigin[i]=m_ArgsInfo.origin_arg[i];
    } else outputOrigin=input->GetOrigin();

    // Set
    resampler->SetSize( outputSize );
    resampler->SetOutputSpacing( outputSpacing );
    resampler->SetOutputOrigin(  outputOrigin );

  }

  if (m_ArgsInfo.verbose_flag) {
    std::cout << "Setting the output size to " << resampler->GetSize() << "..." << std::endl;
    std::cout << "Setting the output spacing to " << resampler->GetOutputSpacing() << "..." << std::endl;
    std::cout << "Setting the output origin to " << resampler->GetOutputOrigin() << "..." << std::endl;
  }

  resampler->SetInput( input );
  resampler->SetTransform( affineTransform );
  resampler->SetInterpolator( genericInterpolator->GetInterpolatorPointer());
  resampler->SetDefaultPixelValue( static_cast<PixelType>(m_ArgsInfo.pad_arg) );

  try {
    resampler->Update();
  } catch(itk::ExceptionObject) {
    std::cerr<<"Error resampling the image"<<std::endl;
  }

  typename OutputImageType::Pointer output = resampler->GetOutput();

  // Output
  typedef itk::ImageFileWriter<OutputImageType> WriterType;
  typename WriterType::Pointer writer = WriterType::New();
  writer->SetFileName(m_ArgsInfo.output_arg);
  writer->SetInput(output);
  writer->Update();

}

//-------------------------------------------------------------------
// Update with the number of dimensions and the pixeltype (components)
//-------------------------------------------------------------------
template<class args_info_type>
template<unsigned int Dimension, class PixelType>
void AffineTransformGenericFilter<args_info_type>::UpdateWithDimAndVectorType()
{
  // ImageTypes
  typedef itk::Image<PixelType, Dimension> InputImageType;
  typedef itk::Image<PixelType, Dimension> OutputImageType;

  // Read the input
  typedef itk::ImageFileReader<InputImageType> InputReaderType;
  typename InputReaderType::Pointer reader = InputReaderType::New();
  reader->SetFileName( m_InputFileName);
  reader->Update();
  typename InputImageType::Pointer input= reader->GetOutput();

  //Filter
  typedef  itk::VectorResampleImageFilter< InputImageType,OutputImageType, double >  ResampleFilterType;
  typename ResampleFilterType::Pointer resampler = ResampleFilterType::New();

  // Matrix
  typename itk::Matrix<double, Dimension+1, Dimension+1> matrix;
  if (m_ArgsInfo.rotate_given || m_ArgsInfo.translate_given)
  {
    if (m_ArgsInfo.matrix_given)
    {
      std::cerr << "You must use either rotate/translate or matrix options" << std::cout;
      return;
    }
    itk::Array<double> transformParameters(2 * Dimension);
    transformParameters.Fill(0.0);
    if (m_ArgsInfo.rotate_given)
    {
      if (Dimension == 2)
        transformParameters[0] = m_ArgsInfo.rotate_arg[0];
      else
        for (unsigned int i = 0; i < 3; i++)
          transformParameters[i] = m_ArgsInfo.rotate_arg[i];
    }
    if (m_ArgsInfo.translate_given)
    {
      int pos = 3;
      if (Dimension == 2)
        pos = 1;
      for (unsigned int i = 0; i < Dimension && i < 3; i++)
        transformParameters[pos++] = m_ArgsInfo.translate_arg[i];
    }
    if (Dimension == 4)
    {
      matrix.SetIdentity();
      itk::Matrix<double, 4, 4> tmp = GetForwardAffineMatrix3D(transformParameters);
      for (unsigned int i = 0; i < 3; ++i)
        for (unsigned int j = 0; j < 3; ++j)
          matrix[i][j] = tmp[i][j];
      for (unsigned int i = 0; i < 3; ++i)
        matrix[i][4] = tmp[i][3];
    }
    else
      matrix = GetForwardAffineMatrix<Dimension>(transformParameters);
  }
  else
  {
    if (m_ArgsInfo.matrix_given)
    {
      matrix= clitk::ReadMatrix<Dimension>(m_ArgsInfo.matrix_arg);
      if (m_Verbose) std::cout << "Reading the matrix..." << std::endl;
    }
    else
      matrix.SetIdentity();
  }
  if (m_Verbose)
    std::cout << "Using the following matrix:" << std::endl
              << matrix << std::endl;
  typename itk::Matrix<double, Dimension, Dimension> rotationMatrix = clitk::GetRotationalPartMatrix(matrix);
  typename itk::Vector<double, Dimension> translationPart = clitk::GetTranslationPartMatrix(matrix);

  // Transform
  typedef itk::AffineTransform<double, Dimension> AffineTransformType;
  typename AffineTransformType::Pointer affineTransform=AffineTransformType::New();
  affineTransform->SetMatrix(rotationMatrix);
  affineTransform->SetTranslation(translationPart);

  // Interp
  typedef clitk::GenericVectorInterpolator<args_info_type, InputImageType, double> GenericInterpolatorType;
  typename GenericInterpolatorType::Pointer genericInterpolator=GenericInterpolatorType::New();
  genericInterpolator->SetArgsInfo(m_ArgsInfo);

  // Properties
  if (m_ArgsInfo.like_given) {
    typename InputReaderType::Pointer likeReader=InputReaderType::New();
    likeReader->SetFileName(m_ArgsInfo.like_arg);
    likeReader->Update();
    resampler->SetSize( likeReader->GetOutput()->GetLargestPossibleRegion().GetSize() );
    resampler->SetOutputSpacing( likeReader->GetOutput()->GetSpacing() );
    resampler->SetOutputOrigin(  likeReader->GetOutput()->GetOrigin() );
  } else {
    //Size
    typename OutputImageType::SizeType outputSize;
    if (m_ArgsInfo.size_given) {
      for(unsigned int i=0; i< Dimension; i++)
        outputSize[i]=m_ArgsInfo.size_arg[i];
    } else outputSize=input->GetLargestPossibleRegion().GetSize();
    std::cout<<"Setting the size to "<<outputSize<<"..."<<std::endl;

    //Spacing
    typename OutputImageType::SpacingType outputSpacing;
    if (m_ArgsInfo.spacing_given) {
      for(unsigned int i=0; i< Dimension; i++)
        outputSpacing[i]=m_ArgsInfo.spacing_arg[i];
    } else outputSpacing=input->GetSpacing();
    std::cout<<"Setting the spacing to "<<outputSpacing<<"..."<<std::endl;

    //Origin
    typename OutputImageType::PointType outputOrigin;
    if (m_ArgsInfo.origin_given) {
      for(unsigned int i=0; i< Dimension; i++)
        outputOrigin[i]=m_ArgsInfo.origin_arg[i];
    } else outputOrigin=input->GetOrigin();
    std::cout<<"Setting the origin to "<<outputOrigin<<"..."<<std::endl;

    // Set
    resampler->SetSize( outputSize );
    resampler->SetOutputSpacing( outputSpacing );
    resampler->SetOutputOrigin(  outputOrigin );

  }

  resampler->SetInput( input );
  resampler->SetTransform( affineTransform );
  resampler->SetInterpolator( genericInterpolator->GetInterpolatorPointer());
  resampler->SetDefaultPixelValue( static_cast<PixelType>(m_ArgsInfo.pad_arg) );

  try {
    resampler->Update();
  } catch(itk::ExceptionObject) {
    std::cerr<<"Error resampling the image"<<std::endl;
  }

  typename OutputImageType::Pointer output = resampler->GetOutput();

  // Output
  typedef itk::ImageFileWriter<OutputImageType> WriterType;
  typename WriterType::Pointer writer = WriterType::New();
  writer->SetFileName(m_ArgsInfo.output_arg);
  writer->SetInput(output);
  writer->Update();

}


} //end clitk

#endif //#define clitkAffineTransformGenericFilter_txx