#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 "itkImage.h"

#include "itkAffineTransform.h"

#include "itkVectorImage.h"

#include "itkImageRegionIteratorWithIndex.h"

#include "itkImageMomentsCalculator.h"

#include "itkTransformFileWriter.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template< class ImageType >

int reflectionMatrixHelper( AntsImage<ImageType> & py_image, unsigned int axis, std::string filename )

{

  typedef typename ImageType::Pointer       ImagePointerType;

  typedef itk::AffineTransform<double, ImageType::ImageDimension> AffineTransformType;

  ImagePointerType image = py_image.ptr;

  typedef typename itk::ImageMomentsCalculator<ImageType> ImageCalculatorType;

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

  calculator->SetImage( image );

  typename ImageCalculatorType::VectorType fixed_center;

  fixed_center.Fill(0);

  calculator->Compute();

  fixed_center = calculator->GetCenterOfGravity();

  typename AffineTransformType::Pointer aff = AffineTransformType::New();

  aff->SetIdentity();

  typename AffineTransformType::ParametersType myoff = aff->GetFixedParameters();

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

    {

    myoff[i] = fixed_center[i];

    }

  typename AffineTransformType::MatrixType mymat = aff->GetMatrix();

  if( axis < ImageType::ImageDimension )

    {

    mymat[axis][axis] = ( -1.0 );

    }

  aff->SetFixedParameters( myoff );

  aff->SetMatrix( mymat );

  typedef itk::TransformFileWriter TransformWriterType;

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

  transformWriter->SetInput( aff );

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

  transformWriter->Update();

  return 1;

}

template <typename ImageType>

int reflectionMatrixEntry( AntsImage<ImageType> & image, unsigned int axis, std::string filename)

{

  return reflectionMatrixHelper<ImageType>( image, axis, filename );

}

void local_reflectionMatrix(nb::module_ &m)

{

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<unsigned char, 2>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<unsigned char, 3>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<unsigned char, 4>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<unsigned int, 2>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<unsigned int, 3>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<unsigned int, 4>>);

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

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

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

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<double, 2>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<double, 3>>);

  m.def("reflectionMatrix", &reflectionMatrixEntry<itk::Image<double, 4>>);

}