#include "antsUtilities.h"
#include "antsUtilities.h"
#include <vnl/vnl_inverse.h>

#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkMatrixOffsetTransformBase.h"
#include "itkTransformFactory.h"
#include "vtkPolyDataReader.h"
#include "itkDisplacementFieldFromMultiTransformFilter.h"
#include "itkTransformFileReader.h"
#include "itkTransformFileWriter.h"

#include "itkLabeledPointSetFileReader.h"
#include "itkLabeledPointSetFileWriter.h"
#include "itkMesh.h"
#include <vtkPolyData.h>
#include <vtkUnstructuredGrid.h>
#include <vtkUnstructuredGridReader.h>
#include <vtkUnstructuredGridWriter.h>
#include <vtkPolyDataReader.h>
#include <vtkPolyDataWriter.h>
#include <vtkPoints.h>

namespace ants
{
vnl_matrix_fixed<double, 4, 4> ConstructNiftiSform(
  vnl_matrix<double> m_dir,
  vnl_vector<double> v_origin,
  vnl_vector<double> v_spacing)
{
  // Set the NIFTI/RAS transform
  vnl_matrix<double>      m_ras_matrix;
  vnl_diag_matrix<double> m_scale, m_lps_to_ras;
  vnl_vector<double>      v_ras_offset;

  // Compute the matrix
  m_scale.set(v_spacing);
  m_lps_to_ras.set(vnl_vector<double>(3, 1.0) );
  m_lps_to_ras[0] = -1;
  m_lps_to_ras[1] = -1;
  m_ras_matrix = m_lps_to_ras * m_dir * m_scale;

  // Compute the vector
  v_ras_offset = m_lps_to_ras * v_origin;

  // Create the larger matrix
  vnl_vector<double> vcol(4, 1.0);
  vcol.update(v_ras_offset);

  vnl_matrix_fixed<double, 4, 4> m_sform;
  m_sform.set_identity();
  m_sform.update(m_ras_matrix);
  m_sform.set_column(3, vcol);
  return m_sform;
}

vnl_matrix_fixed<double, 4, 4> ConstructVTKtoNiftiTransform(
  vnl_matrix<double> m_dir,
  vnl_vector<double> v_origin,
  vnl_vector<double> v_spacing)
{
  vnl_matrix_fixed<double, 4, 4> vox2nii = ConstructNiftiSform(m_dir, v_origin, v_spacing);
  vnl_matrix_fixed<double, 4, 4> vtk2vox;
  vtk2vox.set_identity();
  for( size_t i = 0; i < 3; i++ )
    {
    vtk2vox(i, i) = 1.0 / v_spacing[i];
    vtk2vox(i, 3) = -v_origin[i] / v_spacing[i];
    }
  return vox2nii * vtk2vox;
}

static bool WarpVTKPolyDataMultiTransform_ParseInput(int argc, char * *argv, char *& input_vtk_filename,
                                                     char *& output_vtk_filename,
                                                     char *& reference_image_filename, TRAN_OPT_QUEUE & opt_queue)
{
  opt_queue.clear();
  opt_queue.reserve(argc - 2);

  input_vtk_filename = argv[0];
  output_vtk_filename = argv[1];

  reference_image_filename = NULL;

  int ind = 2;
  while( ind < argc )
    {
    if( strcmp(argv[ind], "-R") == 0 )
      {
      ind++;
      if( ind >= argc )
        {
        return false;
        }
      reference_image_filename = argv[ind];
      }
    else if( strcmp(argv[ind], "-i") == 0 )
      {
      ind++;
      if( ind >= argc )
        {
        return false;
        }
      TRAN_OPT opt;
      opt.filename = argv[ind];
      if( CheckFileType(opt.filename) != AFFINE_FILE )
        {
        std::cout << "file: " << opt.filename
                 << " is not an affine .txt file. Invalid to use '-i' "
                 << std::endl;
        return false;
        }
      opt.file_type = AFFINE_FILE;
      opt.do_affine_inv = true;
      opt_queue.push_back(opt);
      }
    else
      {
      TRAN_OPT opt;
      opt.filename = argv[ind];
      opt.file_type = CheckFileType(opt.filename);
      opt.do_affine_inv = false;
      opt_queue.push_back(opt);
      }
    ind++;
    }

//    if (reference_image_filename == NULL) {
//        std::cout << "the reference image file (-R) must be given!!!"
//        << std::endl;
//        return false;
//    }

  return true;
}

template <int ImageDimension>
void WarpLabeledPointSetFileMultiTransform(char *input_vtk_filename, char *output_vtk_filename,
                                           char *reference_image_filename, TRAN_OPT_QUEUE & opt_queue)
{
  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;
  typedef itk::DisplacementFieldFromMultiTransformFilter<DisplacementFieldType,
                                                         DisplacementFieldType, AffineTransformType> WarperType;
  typedef itk::LinearInterpolateImageFunction<ImageType> FuncType;

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

  typedef itk::ImageFileReader<ImageType> ImageFileReaderType;
  typename ImageFileReaderType::Pointer reader_img =
    ImageFileReaderType::New();
  typename ImageType::Pointer img_ref;

  typename ImageFileReaderType::Pointer reader_img_ref =
    ImageFileReaderType::New();
  if( reference_image_filename )
    {
    reader_img_ref->SetFileName(reference_image_filename);
    reader_img_ref->Update();
    img_ref = reader_img_ref->GetOutput();
    }
  else
    {
    std::cout << "the reference image file (-R) must be given!!!"
             << std::endl;
    return;
    }

  typename WarperType::Pointer warper = WarperType::New();
  // warper->SetInput(img_mov);
  // warper->SetEdgePaddingValue( 0);
  VectorType pad;
  pad.Fill(0);
  // warper->SetEdgePaddingValue(pad);

  typedef itk::TransformFileReader TranReaderType;

  typedef itk::ImageFileReader<DisplacementFieldType>
    FieldReaderType;

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

    switch( opt_queue[i].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_queue[i].do_affine_inv )
          {
          aff->GetInverse(aff);
          }
        // std::cout << aff << std::endl;
        warper->PushBackAffineTransform(aff);
        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();
        // std::cout << field << std::endl;
        warper->PushBackDisplacementFieldTransform(field);
        break;
        }
      default:
        std::cout << "Unknown file type!" << std::endl;
      }
    }

  warper->SetOutputParametersFromImage( img_ref );

  std::cout << "output size: " << warper->GetOutputSize() << std::endl;
  std::cout << "output spacing: " << warper->GetOutputSpacing() << std::endl;

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

  typename DisplacementFieldType::Pointer field_output =
    DisplacementFieldType::New();
  field_output = warper->GetOutput();

  /**
   * Code to read the mesh
   */
  typename ImageType::PointType point;
  typename ImageType::PointType warpedPoint;

  typedef itk::Mesh<float, ImageDimension> MeshType;
  //  typedef itk::LabeledPointSetFileReader<MeshType> VTKReaderType;
  vtkPolyDataReader *vtkreader = vtkPolyDataReader::New();
  vtkreader->SetFileName( input_vtk_filename );
  vtkreader->Update();
  vtkPolyData *mesh = vtkreader->GetOutput();

  /*
  typename VTKReaderType::Pointer vtkreader = VTKReaderType::New();
  vtkreader->SetFileName( input_vtk_filename );
  vtkreader->Update();
  typename MeshType::PointsContainerIterator It
    = vtkreader->GetOutput()->GetPoints()->Begin();

  while( It != vtkreader->GetOutput()->GetPoints()->End() )
    {
    point.CastFrom( It.Value() );
  */
  vnl_matrix_fixed<double, 4, 4> ijk2ras, vtk2ras, lps2ras;
  vtk2ras = ConstructVTKtoNiftiTransform(
      field_output->GetDirection().GetVnlMatrix(),
      field_output->GetOrigin().GetVnlVector(),
      field_output->GetSpacing().GetVnlVector() );
  ijk2ras.set_identity();
  vtk2ras.set_identity();
  lps2ras.set_identity();
  lps2ras(0, 0) = -1; lps2ras(1, 1) = -1;

  // Set up the transforms
  ijk2ras = ConstructNiftiSform(
      field_output->GetDirection().GetVnlMatrix(),
      field_output->GetOrigin().GetVnlVector(),
      field_output->GetSpacing().GetVnlVector() );

  vtk2ras = ConstructVTKtoNiftiTransform(
      field_output->GetDirection().GetVnlMatrix(),
      field_output->GetOrigin().GetVnlVector(),
      field_output->GetSpacing().GetVnlVector() );

  vnl_matrix_fixed<double, 4, 4> ras2ijk = vnl_inverse(ijk2ras);
  vnl_matrix_fixed<double, 4, 4> ras2vtk = vnl_inverse(vtk2ras);
  vnl_matrix_fixed<double, 4, 4> ras2lps = vnl_inverse(lps2ras);
  // Update the coordinates
  for( int k = 0; k < mesh->GetNumberOfPoints(); k++ )
    {
    // Get the point (in whatever format that it's stored)
    vnl_vector_fixed<double, 4> x_mesh, x_ras, x_ijk, v_warp, v_ras;
    vnl_vector_fixed<double, 4> y_ras, y_mesh;
    x_mesh[0] = mesh->GetPoint(k)[0]; x_mesh[1] = mesh->GetPoint(k)[1]; x_mesh[2] = mesh->GetPoint(k)[2];
    x_mesh[3] = 1.0;

    // Map the point into RAS coordinates
    //    if(parm.mesh_coord == RAS)  x_ras = x_mesh;
    //    else if(parm.mesh_coord == LPS) x_ras = lps2ras * x_mesh;
    //    else if(parm.mesh_coord == IJKOS)
    x_ras = vtk2ras * x_mesh;
    //    else   x_ras = ijk2ras * x_mesh;

    // Map the point to IJK coordinates (continuous index)
    x_ijk = ras2ijk * x_ras;

    typename FuncType::ContinuousIndexType ind;
    ind[0] = x_ijk[0];
    ind[1] = x_ijk[1];
    ind[2] = x_ijk[2];
    field_output->TransformContinuousIndexToPhysicalPoint(ind, point);
    //      std::cout << " point " << point << std::endl;
    // std::cout << " point-t " << point << std::endl;
    bool isInside = warper->MultiTransformSinglePoint( point, warpedPoint );
    // if ( isInside ) std::cout << " point-w " << warpedPoint << std::endl;
    if( isInside )
      {
      typename MeshType::PointType newPoint;
      newPoint.CastFrom( warpedPoint );
      //      vtkreader->GetOutput()->SetPoint( It.Index(), newPoint );
      if( ImageDimension == 3 )
        {
        mesh->GetPoints()->SetPoint(k, warpedPoint[0], warpedPoint[1], warpedPoint[2]);
        }
      }
    //    ++It;
    }

  std::string fn = std::string(output_vtk_filename);
  if( fn.rfind(".vtk") == fn.length() - 4 )
    {
    vtkPolyDataWriter *writer = vtkPolyDataWriter::New();
    writer->SetFileName(output_vtk_filename);
    writer->SetInputData(mesh);
    writer->Update();
    }

  /*
  typedef itk::LabeledPointSetFileWriter<MeshType> VTKWriterType;
  typename VTKWriterType::Pointer vtkwriter = VTKWriterType::New();
  vtkwriter->SetFileName( output_vtk_filename );
  vtkwriter->SetInput( vtkreader->GetOutput() );
  vtkwriter->SetMultiComponentScalars( vtkreader->GetMultiComponentScalars() );
  vtkwriter->SetLines( vtkreader->GetLines() );
  vtkwriter->Update();
  */
//    std::string filePrefix = output_vtk_filename;
//    std::string::size_type pos = filePrefix.rfind(".");
//    std::string extension = std::string(filePrefix, pos, filePrefix.length()
//            - 1);
//    filePrefix = std::string(filePrefix, 0, pos);
//
//    std::cout << "output extension is: " << extension << std::endl;
//
//    if (extension != std::string(".mha")) {
//        typedef itk::VectorImageFileWriter<DisplacementFieldType, ImageType>
//        WriterType;
//        typename WriterType::Pointer writer = WriterType::New();
//        writer->SetFileName(output_vtk_filename);
//        writer->SetUseAvantsNamingConvention(true);
//        writer->SetInput(field_output);
//        writer->Update();
//    } else {
//        typedef itk::ImageFileWriter<DisplacementFieldType> WriterType;
//        typename WriterType::Pointer writer = WriterType::New();
//        writer->SetFileName(output_vtk_filename);
//        writer->SetInput(field_output);
//        writer->Update();
//    }
}

template <int ImageDimension>
void ComposeMultiAffine(char * /*input_affine_txt*/, char *output_affine_txt,
                        char *reference_affine_txt, TRAN_OPT_QUEUE & opt_queue)
{
  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;
  // typedef itk::DisplacementFieldFromMultiTransformFilter<DisplacementFieldType,
  // DisplacementFieldType, AffineTransformType> WarperType;

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

  // typedef itk::ImageFileReader<ImageType> ImageFileReaderType;
  // typename ImageFileReaderType::Pointer reader_img = ImageFileReaderType::New();
  // typename ImageType::Pointer img_ref = ImageType::New();

  // typename ImageFileReaderType::Pointer reader_img_ref = ImageFileReaderType::New();

  typename WarperType::Pointer warper = WarperType::New();
  // warper->SetInput(img_mov);
  // warper->SetEdgePaddingValue( 0);
  VectorType pad;
  pad.Fill(0);
  // warper->SetEdgePaddingValue(pad);

  typedef itk::TransformFileReader TranReaderType;

  typedef itk::ImageFileReader<DisplacementFieldType>
    FieldReaderType;

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

    switch( opt_queue[i].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_queue[i].do_affine_inv )
          {
          aff->GetInverse(aff);
          }
        // std::cout << aff << std::endl;
        warper->PushBackAffineTransform(aff);
        cnt_affine++;
        break;
        }
      case DEFORMATION_FILE:
        {
        std::cout << "Compose affine only files: ignore "
                 << opt.filename << std::endl;
        break;
        }
      default:
        std::cout << "Unknown file type!" << std::endl;
      }
    }

  typedef typename AffineTransformType::CenterType PointType;
  PointType aff_center;

  typename AffineTransformType::Pointer aff_ref_tmp;
  if( reference_affine_txt )
    {
    typename TranReaderType::Pointer tran_reader = TranReaderType::New();
    tran_reader->SetFileName(reference_affine_txt);
    tran_reader->Update();
    aff_ref_tmp = dynamic_cast<AffineTransformType *>( (tran_reader->GetTransformList() )->front().GetPointer() );
    }
  else
    {
    if( cnt_affine > 0 )
      {
      std::cout << "the reference affine file for center is selected as the first affine!" << std::endl;
      aff_ref_tmp = ( (warper->GetTransformList() ).begin() )->second.aex.aff;
      }
    else
      {
      std::cout << "No affine input is given. nothing to do ......" << std::endl;
      return;
      }
    }

  aff_center = aff_ref_tmp->GetCenter();
  std::cout << "new center is : " << aff_center << std::endl;

  // warper->PrintTransformList();

  // typename AffineTransformType::Pointer aff_output = warper->ComposeAffineOnlySequence(aff_center);
  typename AffineTransformType::Pointer aff_output = AffineTransformType::New();
  warper->ComposeAffineOnlySequence(aff_center, aff_output);
  typedef itk::TransformFileWriter TranWriterType;
  typename TranWriterType::Pointer tran_writer = TranWriterType::New();
  tran_writer->SetFileName(output_affine_txt);
  tran_writer->SetInput(aff_output);
  tran_writer->Update();

  std::cout << "wrote file to : " << output_affine_txt << std::endl;
}

// entry point for the library; parameter 'args' is equivalent to 'argv' in (argc,argv) of commandline parameters to
// 'main()'
int WarpVTKPolyDataMultiTransform( std::vector<std::string> args, std::ostream* )
{
  // 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(), "WarpVTKPolyDataMultiTransform" );

  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] = 0;
  // 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 <= 4 )
    {
    std::cout
      << "WarpLabeledPointSetFileMultiTransform ImageDimension inputVTKFile "
      << "outputVTKFile [-R reference_image] "
      << "{[deformation_field | [-i] affine_transform_txt ]}"
      << std::endl;
    return EXIT_FAILURE;
    }

  TRAN_OPT_QUEUE opt_queue;
  //    char *moving_image_filename = NULL;
  char *input_vtk_filename = NULL;
  char *output_vtk_filename = NULL;
  char *reference_image_filename = NULL;

  int  kImageDim = atoi(argv[1]);

  const bool is_parsing_ok = WarpVTKPolyDataMultiTransform_ParseInput(argc - 2, argv + 2, input_vtk_filename, output_vtk_filename,
                                                           reference_image_filename, opt_queue);

  if( is_parsing_ok )
    {
    switch( CheckFileType(output_vtk_filename) )
      {
      case DEFORMATION_FILE:
        {
        if( reference_image_filename == NULL )
          {
          std::cout << "the reference image file (-R) must be given!!!"
                   << std::endl;
          return false;
          }

        std::cout << "output_vtk_filename: " << output_vtk_filename
                 << std::endl;
        std::cout << "reference_image_filename: ";
        if( reference_image_filename )
          {
          std::cout << reference_image_filename << std::endl;
          }
        else
          {
          std::cout << "NULL" << std::endl;
          }
        DisplayOptQueue(opt_queue);

        switch( kImageDim )
          {
          case 2:
            {
            WarpLabeledPointSetFileMultiTransform<2>(input_vtk_filename,
                                                     output_vtk_filename,
                                                     reference_image_filename, opt_queue);
            break;
            }
          case 3:
            {
            WarpLabeledPointSetFileMultiTransform<3>(input_vtk_filename,
                                                     output_vtk_filename,
                                                     reference_image_filename, opt_queue);
            break;
            }
          }
        break;
        }

      case AFFINE_FILE:
        {
        std::cout << "output_affine_txt: " << output_vtk_filename
                 << std::endl;
        std::cout << "reference_affine_txt: ";
        if( reference_image_filename )
          {
          std::cout << reference_image_filename << std::endl;
          }
        else
          {
          std::cout << "NULL" << std::endl;
          }
        DisplayOptQueue(opt_queue);

        switch( kImageDim )
          {
          case 2:
            {
            ComposeMultiAffine<2>(input_vtk_filename, output_vtk_filename,
                                  reference_image_filename, opt_queue);
            break;
            }
          case 3:
            {
            ComposeMultiAffine<3>(input_vtk_filename, output_vtk_filename,
                                  reference_image_filename, opt_queue);
            break;
            }
          }
        break;
        }

      default:
        std::cout << "Unknow output file format: " << output_vtk_filename << std::endl;
        break;
      }
    }
  else
    {
    std::cout << "Input error!" << std::endl;
    }

  return EXIT_FAILURE;
}
} // namespace ants
