#include <nanobind/nanobind.h>

#include <nanobind/stl/vector.h>

#include <nanobind/stl/string.h>

#include <nanobind/stl/tuple.h>

#include <nanobind/stl/list.h>

#include <nanobind/ndarray.h>

#include <nanobind/stl/shared_ptr.h>

#include <exception>

#include <vector>

#include <string>

#include "itkAffineTransform.h"

#include "itkImage.h"

#include "itkImageFileWriter.h"

#include "itkImageMomentsCalculator.h"

#include "itkResampleImageFilter.h"

#include "itkTransformFileWriter.h"

#include "vnl/vnl_inverse.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

/*

template< unsigned int ImageDimension >

int antsReoHelper(

  typename itk::Image< float , ImageDimension >::Pointer image1,

  std::string txfn, std::vector<int> axis , std::vector<int> axis2,

  std::vector<int> doReflection, std::vector<int> doScale )

{

  typedef double RealType;

  typedef itk::Image< float , ImageDimension > ImageType;

  typedef typename itk::ImageMomentsCalculator<ImageType> ImageCalculatorType;

  typedef itk::AffineTransform<RealType, ImageDimension> AffineType;

  typedef typename ImageCalculatorType::MatrixType       MatrixType;

  typedef itk::Vector<float, ImageDimension>  VectorType;

  VectorType ccg1;

  VectorType cpm1;

  MatrixType cpa1;

  VectorType ccg2;

  VectorType cpm2;

  MatrixType cpa2;

  typename ImageCalculatorType::Pointer calculator1 =

    ImageCalculatorType::New();

  calculator1->SetImage(  image1 );

  typename ImageCalculatorType::VectorType fixed_center;

  fixed_center.Fill(0);

  // was a try-catch block around this in ANTsR

  calculator1->Compute();

  fixed_center = calculator1->GetCenterOfGravity();

  ccg1 = calculator1->GetCenterOfGravity();

  cpm1 = calculator1->GetPrincipalMoments();

  cpa1 = calculator1->GetPrincipalAxes();

  unsigned int eigind1 = 1;

  unsigned int eigind2 = 1;

  typedef vnl_vector<RealType> EVectorType;

  typedef vnl_matrix<RealType> EMatrixType;

  EVectorType evec1_2ndary = cpa1.GetVnlMatrix().get_row( eigind2 );

  EVectorType evec1_primary = cpa1.GetVnlMatrix().get_row( eigind1 );

  EVectorType evec2_2ndary;

  evec2_2ndary.set_size( ImageDimension );

  evec2_2ndary.fill(0);

  EVectorType evec2_primary;

  evec2_primary.set_size( ImageDimension );

  evec2_primary.fill(0);

  for ( unsigned int i = 0; i < ImageDimension; i++ )

    evec2_primary[i] = axis[i];

  // Solve Wahba's problem http://en.wikipedia.org/wiki/Wahba%27s_problem

  EMatrixType B = outer_product( evec2_primary, evec1_primary );

  if( ImageDimension == 3 )

    {

    for ( unsigned int i = 0; i < ImageDimension; i++ )

        evec2_primary[i] = axis2[i];

    B = outer_product( evec2_2ndary, evec1_2ndary )

      + outer_product( evec2_primary, evec1_primary );

    }

    vnl_svd<RealType>    wahba( B );

    vnl_matrix<RealType> A_solution = wahba.V() * wahba.U().transpose();

    A_solution = vnl_inverse( A_solution );

    RealType det = vnl_determinant( A_solution  );

    if( det < 0 )

      {

      vnl_matrix<RealType> id( A_solution );

      id.set_identity();

      for( unsigned int i = 0; i < ImageDimension; i++ )

        {

        if( A_solution( i, i ) < 0 )

          {

          id( i, i ) = -1.0;

          }

        }

        A_solution =  A_solution * id.transpose();

      }

    if ( doReflection[0] == 1 ||  doReflection[0] == 3 )

      {

      vnl_matrix<RealType> id( A_solution );

      id.set_identity();

      id = id - 2.0 * outer_product( evec2_primary , evec2_primary  );

      A_solution = A_solution * id;

      }

    if ( doReflection[0] > 1 )

      {

      vnl_matrix<RealType> id( A_solution );

      id.set_identity();

      id = id - 2.0 * outer_product( evec1_primary , evec1_primary  );

      A_solution = A_solution * id;

      }

    if ( doScale[0] > 0 )

      {

      vnl_matrix<RealType> id( A_solution );

      id.set_identity();

      id = id * doScale[0];

      A_solution = A_solution * id;

      }

    det = vnl_determinant( A_solution  );

    std::cout << " det " << det << std::endl;

    std::cout << " A_solution " << std::endl;

    std::cout << A_solution << std::endl;

    typename AffineType::Pointer affine1 = AffineType::New();

    typename AffineType::OffsetType trans = affine1->GetOffset();

    itk::Point<RealType, ImageDimension> trans2;

    trans2.Fill(0);

    for( unsigned int i = 0; i < ImageDimension; i++ )

      {

      trans2[i] =  fixed_center[i] * ( 1 );

      }

    affine1->SetIdentity();

    affine1->SetOffset( trans );

    affine1->SetMatrix( A_solution );

    affine1->SetCenter( trans2 );

    // write tx

    typedef itk::TransformFileWriter TransformWriterType;

    typename TransformWriterType::Pointer transformWriter = TransformWriterType::New();

    transformWriter->SetInput( affine1 );

    transformWriter->SetFileName( txfn.c_str() );

    transformWriter->Update();

  return 0;

}

template <typename ImageType, unsigned int Dimension>

int reorientImage(  py::capsule in_image, std::string txfn,

                      std::vector<int> axis1, std::vector<int> axis2,

                      std::vector<int> rrfl, std::vector<int> rscl )

{

  typedef typename ImageType::Pointer ImagePointerType;

  ImagePointerType itk_image = as<ImageType>( in_image );

  antsReoHelper<Dimension>( itk_image, txfn, axis1, axis2, rrfl, rscl );

  return 0;

}

*/

template <typename ImageType, unsigned int Dimension>

std::vector<double> centerOfMass( AntsImage<ImageType> & image  )

{

  typedef typename ImageType::Pointer ImagePointerType;

  ImagePointerType itkimage = image.ptr;

  typedef typename itk::ImageMomentsCalculator<ImageType> ImageCalculatorType;

  typename ImageCalculatorType::VectorType com( Dimension );

  com.Fill( 0 );

  std::vector<double> myCoM( Dimension );

  typename ImageCalculatorType::Pointer calculator1 = ImageCalculatorType::New();

  calculator1->SetImage( itkimage );

  calculator1->Compute();

  com = calculator1->GetCenterOfGravity();

  for ( unsigned int k = 0; k < Dimension; k++ )

  {

    myCoM[ k ] = com[ k ];

  }

  return myCoM;

}

void local_reorientImage(nb::module_ &m)

{

//  m.def("reorientImageF2", &reorientImage<itk::Image<float, 2>, 2>);

//  m.def("reorientImageF3", &reorientImage<itk::Image<float, 3>, 3>);

//  m.def("reorientImageF4", &reorientImage<itk::Image<float, 4>, 4>);

  m.def("centerOfMass", &centerOfMass<itk::Image<float, 2>, 2>);

  m.def("centerOfMass", &centerOfMass<itk::Image<float, 3>, 3>);

  m.def("centerOfMass", &centerOfMass<itk::Image<float, 4>, 4>);

}