
#include "antsUtilities.h"
#include "antsAllocImage.h"
#include "itkImageFileReader.h"
#include "itkVariableLengthVector.h"

#include "itkImageFileWriter.h"
#include "itkMatrixOffsetTransformBase.h"
#include "itkTransformFactory.h"
#include "itkWarpTensorImageMultiTransformFilter.h"
#include "itkTransformFileReader.h"
#include "itkVectorNearestNeighborInterpolateImageFunction.h"
#include "ReadWriteData.h"
#include "itkWarpImageMultiTransformFilter.h"
#include "itkExpTensorImageFilter.h"

namespace ants
{
static bool WarpTensorImageMultiTransform_ParseInput(int argc, char * *argv, char *& moving_image_filename,
                                                     char *& output_image_filename,
                                                     TRAN_OPT_QUEUE & opt_queue, MISC_OPT & misc_opt)
{
  opt_queue.clear();
  opt_queue.reserve(argc - 2);

  misc_opt.reference_image_filename = ITK_NULLPTR;
  misc_opt.use_NN_interpolator = false;
  misc_opt.use_TightestBoundingBox = false;
  misc_opt.use_RotationHeader = false;

  moving_image_filename = argv[0];
  output_image_filename = argv[1];

  int  ind = 2;
  bool set_current_affine_inv = false;

  while( ind < argc )
    {
    if( strcmp(argv[ind], "--use-NN") == 0 )
      {
      misc_opt.use_NN_interpolator = true;
      }
    else if( strcmp(argv[ind], "-R") == 0 )
      {
      ind++; if( ind >= argc )
        {
        return false;
        }
      misc_opt.reference_image_filename = argv[ind];
      }
    else if( (strcmp(argv[ind], "--tightest-bounding-box") == 0) &&  (strcmp(argv[ind], "-R") != 0)  )
      {
      misc_opt.use_TightestBoundingBox = true;
      }
    else if( strcmp(argv[ind], "--reslice-by-header") == 0 )
      {
      misc_opt.use_RotationHeader = true;
      TRAN_OPT opt;
      opt.file_type = IMAGE_AFFINE_HEADER;
      opt.do_affine_inv = false;
      opt_queue.push_back(opt);
      }
    else if( strcmp(argv[ind], "--Id") == 0 )
      {
      TRAN_OPT opt;
      opt.filename = "--Id";
      opt.do_affine_inv = false;
      opt.file_type = IDENTITY_TRANSFORM;
      opt_queue.push_back(opt);
      }
    else if( strcmp(argv[ind], "--moving-image-header") == 0 || strcmp(argv[ind], "-mh") == 0 )
      {
      TRAN_OPT opt;
      opt.file_type = IMAGE_AFFINE_HEADER;
      opt.filename = moving_image_filename;
      //            opt.do_affine_inv = false;
      SetAffineInvFlag(opt, set_current_affine_inv);
      opt_queue.push_back(opt);
      }
    else if( strcmp(argv[ind], "--reference-image-header") == 0 || strcmp(argv[ind], "-rh") == 0 )
      {
      if( misc_opt.reference_image_filename == ITK_NULLPTR )
        {
        std::cout
          << "reference image filename is not given yet. Specify it with -R before --reference-image-header / -rh."
          << std::endl;
        return false;
        }

      TRAN_OPT opt;
      opt.file_type = IMAGE_AFFINE_HEADER;
      opt.filename = misc_opt.reference_image_filename;
      //            opt.do_affine_inv = false;
      SetAffineInvFlag(opt, set_current_affine_inv);
      opt_queue.push_back(opt);
      }
    else if( strcmp(argv[ind], "-i") == 0 )
      {
      set_current_affine_inv = true;
      }

    else if( strcmp(argv[ind], "--ANTS-prefix") == 0 )
      {
      ind++;
      std::string prefix = argv[ind];
      std::string path, name, ext;
      FilePartsWithgz(prefix, path, name, ext);
      if( ext == "" )
        {
        ext = ".nii.gz";
        }

      std::string deform_file_name, x_deform_name;
      deform_file_name = path + name + std::string("Warp") + ext;
      x_deform_name = path + name + std::string("Warpxvec") + ext;
      if( CheckFileExistence(x_deform_name.c_str() ) )
        {
        TRAN_OPT opt;
        opt.filename = deform_file_name.c_str();
        opt.file_type = CheckFileType(opt.filename.c_str() );
        opt.do_affine_inv = false;
        opt_queue.push_back(opt);
        std::cout << "found deformation file: " << opt.filename << std::endl;
        DisplayOpt(opt);
        }

      std::string affine_file_name;
      affine_file_name = path + name + std::string("Affine.txt");
      if( CheckFileExistence(affine_file_name.c_str() ) )
        {
        TRAN_OPT opt;
        opt.filename = affine_file_name.c_str();
        opt.file_type = CheckFileType(opt.filename.c_str() );
        opt.do_affine_inv = false;
        opt_queue.push_back(opt);
        std::cout << "found affine file: " << opt.filename << std::endl;
        DisplayOpt(opt);
        }
      }
    else if( strcmp(argv[ind], "--ANTS-prefix-invert") == 0 )
      {
      ind++;
      std::string prefix = argv[ind];
      std::string path, name, ext;
      FilePartsWithgz(prefix, path, name, ext);
      if( ext == "" )
        {
        ext = ".nii.gz";
        }

      std::string affine_file_name;
      affine_file_name = path + name + std::string("Affine.txt");
      if( CheckFileExistence(affine_file_name.c_str() ) )
        {
        TRAN_OPT opt;
        opt.filename = affine_file_name.c_str();
        opt.file_type = CheckFileType(opt.filename.c_str() );
        opt.do_affine_inv = true;
        opt_queue.push_back(opt);
        std::cout << "found affine file: " << opt.filename << std::endl;
        DisplayOpt(opt);
        }

      std::string deform_file_name, x_deform_name;
      deform_file_name = path + name + std::string("InverseWarp.nii.gz");
      x_deform_name = path + name + std::string("InverseWarpxvec.nii.gz");
      if( CheckFileExistence(x_deform_name.c_str() ) )
        {
        TRAN_OPT opt;
        opt.filename = deform_file_name.c_str();
        opt.file_type = CheckFileType(opt.filename.c_str() );
        opt.do_affine_inv = false;
        opt_queue.push_back(opt);
        std::cout << "found deformation file: " << opt.filename << std::endl;
        DisplayOpt(opt);
        }
      }
    else
      {
      TRAN_OPT opt;
      opt.filename = argv[ind];
      opt.file_type = CheckFileType(opt.filename.c_str() );
      opt.do_affine_inv = false;
      if( opt.file_type == AFFINE_FILE )
        {
        SetAffineInvFlag(opt, set_current_affine_inv);
        }
      else if( opt.file_type == DEFORMATION_FILE && set_current_affine_inv )
        {
        std::cout << "Ignore inversion of non-affine file type! " << std::endl;
        std::cout << "opt.do_affine_inv:" << opt.do_affine_inv << std::endl;
        }

      opt_queue.push_back(opt);
      DisplayOpt(opt);
      }
    ind++;
    }

  if( misc_opt.use_RotationHeader )
    {
    //                if (misc_opt.reference_image_filename) {
    //                    opt_queue[0].filename = misc_opt.reference_image_filename;
    //                } else {
    opt_queue[0].filename = "--Id";
    opt_queue[0].file_type = IDENTITY_TRANSFORM;
    opt_queue[0].do_affine_inv = false;
    //                }

    //               TRAN_OPT opt;
    //               opt.file_type = IMAGE_AFFINE_HEADER;
    //               opt.filename = moving_image_filename;
    //               opt.do_affine_inv = true;
    //               opt_queue.push_back(opt);
    //
    //               std::cout << "Use Rotation Header!" << std::endl;
    }

  return true;
}

template <class TAffineTransform>
static void GetIdentityTransform(typename TAffineTransform::Pointer & aff)
{
  aff = TAffineTransform::New();
  aff->SetIdentity();
}

#if 0
template <typename TensorImageType, typename ImageType>
static void
DirectionCorrect( typename TensorImageType::Pointer img_mov, typename ImageType::Pointer img_ref )
{
  itk::ImageRegionIteratorWithIndex<TensorImageType> it(img_mov, img_mov->GetLargestPossibleRegion() );

  typename TensorImageType::DirectionType::InternalMatrixType direction = img_mov->GetDirection().GetTranspose()
    * img_ref->GetDirection().GetVnlMatrix();

  if( !direction.is_identity( 0.00001 ) )
    {
    while( !it.IsAtEnd() )
      {
      typename TensorImageType::DirectionType::InternalMatrixType  dt;
      dt(0, 0) = it.Value()[0];
      dt(0, 1) = dt(1, 0) = it.Value()[1];
      dt(0, 2) = dt(2, 0) = it.Value()[2];
      dt(1, 1) = it.Value()[3];
      dt(1, 2) = dt(2, 1) = it.Value()[4];
      dt(2, 2) = it.Value()[5];

      dt = direction * dt * direction.transpose();

      typename TensorImageType::PixelType outDt;

      outDt[0] = dt(0, 0);
      outDt[1] = dt(0, 1);
      outDt[2] = dt(0, 2);
      outDt[3] = dt(1, 1);
      outDt[4] = dt(1, 2);
      outDt[5] = dt(2, 2);

      it.Set( outDt );

      ++it;
      }
    }
}

#endif

template <int ImageDimension>
static void WarpImageMultiTransform(char *moving_image_filename, char *output_image_filename,
                                    TRAN_OPT_QUEUE & opt_queue, MISC_OPT & misc_opt)
{
  // typedef itk::Vector<float,6> PixelType;
  typedef itk::SymmetricSecondRankTensor<float,
                                         3>                       PixelType;
  typedef itk::Image<PixelType,
                     ImageDimension>                              TensorImageType;
  typedef itk::Image<float,
                     ImageDimension>                              ImageType;
  typedef itk::Vector<float,
                      ImageDimension>                             VectorType;
  typedef itk::Image<VectorType,
                     ImageDimension>                              DisplacementFieldType;
  typedef itk::MatrixOffsetTransformBase<double, ImageDimension,
                                         ImageDimension>          AffineTransformType;
  typedef itk::WarpImageMultiTransformFilter<ImageType, ImageType, DisplacementFieldType,
                                             AffineTransformType> WarperType;

  itk::TransformFactory<AffineTransformType>::RegisterTransform();

  typedef itk::ImageFileReader<ImageType> ImageFileReaderType;
  // typename ImageFileReaderType::Pointer reader_img = ImageFileReaderType::New();
  // reader_img->SetFileName(moving_image_filename);
  // reader_img->Update();
  // typename ImageType::Pointer img_mov = ImageType::New();
  // img_mov = reader_img->GetOutput();
  typename TensorImageType::Pointer img_mov;
  ReadTensorImage<TensorImageType>(img_mov, moving_image_filename, true);

  typename ImageType::Pointer img_ref;

  typename ImageFileReaderType::Pointer reader_img_ref = ImageFileReaderType::New();
  if( misc_opt.reference_image_filename )
    {
    reader_img_ref->SetFileName(misc_opt.reference_image_filename);
    reader_img_ref->Update();
    img_ref = reader_img_ref->GetOutput();
    }

  // Convert to reference image tensor basis
  // DirectionCorrect<TensorImageType, ImageType>(img_mov, img_ref);

  typename TensorImageType::Pointer img_output =
    AllocImage<TensorImageType>(img_ref);
  for( unsigned int tensdim = 0;  tensdim < 6;  tensdim++ )
    {
    typedef itk::VectorIndexSelectionCastImageFilter<TensorImageType, ImageType> IndexSelectCasterType;
    typename IndexSelectCasterType::Pointer fieldCaster = IndexSelectCasterType::New();
    fieldCaster->SetInput( img_mov );
    fieldCaster->SetIndex( tensdim );
    fieldCaster->Update();
    typename ImageType::Pointer tenscomponent = fieldCaster->GetOutput();
    tenscomponent->SetSpacing(img_mov->GetSpacing() );
    tenscomponent->SetOrigin(img_mov->GetOrigin() );
    tenscomponent->SetDirection(img_mov->GetDirection() );

    typename WarperType::Pointer  warper = WarperType::New();
    warper->SetInput(tenscomponent);
    //      PixelType nullPix;
    // nullPix.Fill(0);
    warper->SetEdgePaddingValue(0);

    if( misc_opt.use_NN_interpolator )
      {
      typedef typename itk::NearestNeighborInterpolateImageFunction<ImageType,
                                                                    typename WarperType::CoordRepType>
        NNInterpolateType;
      typename NNInterpolateType::Pointer interpolator_NN = NNInterpolateType::New();
      std::cout << "Haha" << std::endl;
      warper->SetInterpolator(interpolator_NN);
      }

    typedef itk::TransformFileReader                    TranReaderType;
    typedef itk::ImageFileReader<DisplacementFieldType> FieldReaderType;

    unsigned int transcount = 0;
    const int    kOptQueueSize = opt_queue.size();
    for( int i = 0; i < kOptQueueSize; i++ )
      {
      const TRAN_OPT & opt = opt_queue[i];

      switch( opt.file_type )
        {
        case AFFINE_FILE:
          {
          typename TranReaderType::Pointer tran_reader = TranReaderType::New();
          tran_reader->SetFileName(opt.filename);
          tran_reader->Update();
          typename AffineTransformType::Pointer aff = dynamic_cast<AffineTransformType *>
            ( (tran_reader->GetTransformList() )->front().GetPointer() );
          if( opt.do_affine_inv )
            {
            typename AffineTransformType::Pointer aff_inv = AffineTransformType::New();
            aff->GetInverse(aff_inv);
            aff = aff_inv;
            }
          // std::cout <<" aff " << transcount <<  std::endl;
          warper->PushBackAffineTransform(aff);
          if( transcount == 0 )
            {
            warper->SetOutputParametersFromImage( img_mov );
            }
          transcount++;
          }
          break;
        case IDENTITY_TRANSFORM:
          {
          typename AffineTransformType::Pointer aff;
          GetIdentityTransform<AffineTransformType>(aff);
          // std::cout << " aff id" << transcount << std::endl;
          warper->PushBackAffineTransform(aff);
          transcount++;
          }
          break;
        case IMAGE_AFFINE_HEADER:
          {
          typename AffineTransformType::Pointer aff = AffineTransformType::New();
          typename ImageFileReaderType::Pointer reader_image_affine = ImageFileReaderType::New();
          reader_image_affine->SetFileName(opt.filename);
          reader_image_affine->Update();
          typename ImageType::Pointer img_affine = reader_image_affine->GetOutput();

          GetAffineTransformFromImage<ImageType, AffineTransformType>(img_affine, aff);

          if( opt.do_affine_inv )
            {
            typename AffineTransformType::Pointer aff_inv = AffineTransformType::New();
            aff->GetInverse(aff_inv);
            aff = aff_inv;
            }

          // std::cout <<" aff from image header " << transcount <<  std::endl;
          warper->PushBackAffineTransform(aff);

          //            if (transcount==0){
          //                warper->SetOutputParametersFromImage( img_mov );
          //            }

          transcount++;
          }
          break;
        case DEFORMATION_FILE:
          {
          typename FieldReaderType::Pointer field_reader = FieldReaderType::New();
          field_reader->SetFileName( opt.filename );
          field_reader->Update();
          typename DisplacementFieldType::Pointer field = field_reader->GetOutput();
          warper->PushBackDisplacementFieldTransform(field);
          warper->SetOutputParametersFromImage( field );
          transcount++;
          }
          break;
        default:
          {
          std::cout << "Unknown file type!" << std::endl;
          }
        }
      }

    // warper->PrintTransformList();

    if( img_ref.IsNotNull() )
      {
      warper->SetOutputParametersFromImage( img_ref );
      }
    else
      {
      if( misc_opt.use_TightestBoundingBox == true )
        {
        // compute the desired spacking after inputting all the transform files using the

        typename ImageType::SizeType largest_size;
        typename ImageType::PointType origin_warped;
        GetLargestSizeAfterWarp<WarperType, TensorImageType>(warper, img_mov, largest_size, origin_warped);
        warper->SetOutputParametersFromImage( img_mov );
        warper->SetOutputSize(largest_size);
        warper->SetOutputOrigin(origin_warped);
          {
          typename ImageType::DirectionType d;
          d.SetIdentity();
          warper->SetOutputDirection(d);
          }
        }
      }

    // std::cout << "output origin: " << warper->GetOutputOrigin() << std::endl;
    // std::cout << "output size: " << warper->GetOutputSize() << std::endl;
    // std::cout << "output spacing: " << warper->GetOutputSpacing() << std::endl;
    //    std::cout << "output direction: " << warper->GetOutputDirection() << std::endl;

    // warper->PrintTransformList();
    warper->DetermineFirstDeformNoInterp();
    warper->Update();

    typedef itk::ImageRegionIteratorWithIndex<TensorImageType> Iterator;
    Iterator vfIter2( img_output, img_output->GetLargestPossibleRegion() );
    for(  vfIter2.GoToBegin(); !vfIter2.IsAtEnd(); ++vfIter2 )
      {
      PixelType tens = vfIter2.Get();
      tens[tensdim] = warper->GetOutput()->GetPixel(vfIter2.GetIndex() );
      vfIter2.Set(tens);
      }
    }

  // DirectionCorrect<TensorImageType>(img_output, img_mov);

  WriteTensorImage<TensorImageType>(img_output, output_image_filename, true);
}

// entry point for the library; parameter 'args' is equivalent to 'argv' in (argc,argv) of commandline parameters to
// 'main()'
int WarpTensorImageMultiTransform( std::vector<std::string> args, std::ostream* /*out_stream = NULL */ )
{
  // put the arguments coming in as 'args' into standard (argc,argv) format;
  // 'args' doesn't have the command name as first, argument, so add it manually;
  // 'args' may have adjacent arguments concatenated into one argument,
  // which the parser should handle
  args.insert( args.begin(), "WarpTensorImageMultiTransform" );

  int     argc = args.size();
  char* * argv = new char *[args.size() + 1];
  for( unsigned int i = 0; i < args.size(); ++i )
    {
    // allocate space for the string plus a null character
    argv[i] = new char[args[i].length() + 1];
    std::strncpy( argv[i], args[i].c_str(), args[i].length() );
    // place the null character in the end
    argv[i][args[i].length()] = '\0';
    }
  argv[argc] = ITK_NULLPTR;
  // class to automatically cleanup argv upon destruction
  class Cleanup_argv
  {
public:
    Cleanup_argv( char* * argv_, int argc_plus_one_ ) : argv( argv_ ), argc_plus_one( argc_plus_one_ )
    {
    }

    ~Cleanup_argv()
    {
      for( unsigned int i = 0; i < argc_plus_one; ++i )
        {
        delete[] argv[i];
        }
      delete[] argv;
    }

private:
    char* *      argv;
    unsigned int argc_plus_one;
  };
  Cleanup_argv cleanup_argv( argv, argc + 1 );

  // antscout->set_stream( out_stream );

  if( argc <= 3 )
    {
    std::cout
      <<
      "WarpImageMultiTransform ImageDimension moving_image output_image [-R reference_image | --tightest-bounding-box] (--reslice-by-header) [--use-NN (use Nearest Neighbor Interpolator)]"
      <<
      "[--ANTS-prefix prefix-name | --ANTS-prefix-invert prefix-name] {[deformation_field | [-i] affine_transform_txt | --Id | [-i] --moving-image-header / -mh  | [-i] --reference-image-header / -rh]}"
      << std::endl
      << "Example:" << std::endl
      <<
      "Reslice the image: WarpImageMultiTransform 3 Imov.nii Iout.nii --tightest-bounding-box --reslice-by-header"
      << std::endl
      <<
      "Reslice the image to a reference image: WarpImageMultiTransform 3 Imov.nii Iout.nii -R Iref.nii --tightest-bounding-box --reslice-by-header"
      << std::endl
      << "Note:" << std::endl
      << "-i will use the inversion of the following affine transform." << std::endl
      <<
      "--tightest-bounding-box will be overrided by -R reference_image if given. It computes the tightest bounding box using all the affine transformations."
      << std::endl
      << "--Id uses the identity transform." << std::endl
      <<
      "--moving-image-header or -mh in short will use the orientation header of the moving image file. This is typically not used with --reslice-by-header."
      << std::endl
      <<
      "--reference-image-header or -rh in short will use the orientation header of the fixed image file. This is typically not used with --reslice-by-header."
      << std::endl
      <<
      "--reslice-by-header uses the orientation matrix and origin encoded in the image file header. It can be used together with -R. "
      << "This is typically not used together with any other transforms. "
      << "--reslice-by-header is equvalient to -i -mh, or -fh -i -mh if used together with -R. " << std::endl;
    std::cout << std::endl
             << "For ANTS users:" << std::endl
             << "To use with the deformation field and the affine transform files generated from ANTS:" << std::endl
             << "--ANTS-prefix prefix-name" << std::endl
             << "--ANTS-prefix-invert prefix-name" << std::endl
             << "Example:" << std::endl
             << "3 moving_image output_image -R reference_image --ANTS-prefix abcd.nii.gz" << std::endl
             <<
      "Applies abcdWarpxvec.nii.gz/abcdWarpyvec.nii.gz/abcdWarpzvec.nii.gz and then abcdAffine.txt. Use this with ANTS to get the moving_image warped into the reference_image domain. "
             << std::endl
             << "3 reference_image output_image -R moving_image --ANTS-prefix-invert abcd.nii.gz --ANTS-invert"
             << std::endl
             <<
      "Applies the inversion of abcdAffine.txt and then abcdInverseWarpxvec.nii.gz/abcdInverseWarpyvec.nii.gz/abcdInverseWarpzvec.nii.gz. Use this with ANTS to get the reference_image warped into the moving_image domain. "
             << std::endl
             << "Note: " << std::endl
             << "prefix name \"abcd\" without any extension will use \".nii.gz\" by default" << std::endl;
    if( argc >= 2 &&
        ( std::string( argv[1] ) == std::string("--help") || std::string( argv[1] ) == std::string("-h") ) )
      {
      return EXIT_SUCCESS;
      }
    return EXIT_FAILURE;
    }

  TRAN_OPT_QUEUE opt_queue;
  char *         moving_image_filename = ITK_NULLPTR;
  char *         output_image_filename = ITK_NULLPTR;

  MISC_OPT misc_opt;

  const int  kImageDim = atoi(argv[1]);
  const bool is_parsing_ok =
    WarpTensorImageMultiTransform_ParseInput(argc - 2, argv + 2, moving_image_filename, output_image_filename,
                                             opt_queue,
                                             misc_opt);

  if( is_parsing_ok )
    {
    std::cout << "moving_image_filename: " << moving_image_filename << std::endl;
    std::cout << "output_image_filename: " << output_image_filename << std::endl;
    std::cout << "reference_image_filename: ";
    if( misc_opt.reference_image_filename )
      {
      std::cout << misc_opt.reference_image_filename << std::endl;
      }
    else
      {
      std::cout << "NULL" << std::endl;
      }
    DisplayOptQueue(opt_queue);

    switch( kImageDim )
      {
      case 2:
        {
        WarpImageMultiTransform<2>(moving_image_filename, output_image_filename, opt_queue, misc_opt);
        }
        break;
      case 3:
        {
        WarpImageMultiTransform<3>(moving_image_filename, output_image_filename, opt_queue, misc_opt);
        }
        break;
      }
    }
  else
    {
    std::cout << "Input error!" << std::endl;
    }
  return EXIT_FAILURE;
}
} // namespace ants
