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

#include "itkImageFileReader.h"

#include "itkVectorCurvatureAnisotropicDiffusionImageFilter.h"

#include "itkVectorToRGBImageAdaptor.h"

#include "itkRGBToVectorImageAdaptor.h"

#include "itkCastImageFilter.h"

#include "mergeChannels.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template <typename VectorImageType, typename ImageOfVectorsType, typename PixelType, unsigned int Dimension>

typename VectorImageType::Pointer imageOfVectors_to_VectorImage( typename ImageOfVectorsType::Pointer imageOfVectors )

{

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  typedef typename ImageOfVectorsType::Pointer ImageOfVectorsPointerType;

  typedef itk::Image<PixelType, Dimension> ImageType;

  // split channels

  std::vector<AntsImage<ImageType>> myChannels;

  AntsImage<ImageOfVectorsType> myCapsule = { imageOfVectors };

  myChannels = splitChannels<ImageOfVectorsType,ImageType>( myCapsule );

  // now merge channels into a vectorimage

  VectorImagePointerType finalImage = mergeChannels2<ImageType, VectorImageType>( myChannels );

  return finalImage;

}

template <typename VectorImageType, typename ImageOfVectorsType, typename PixelType, unsigned int Dimension>

typename ImageOfVectorsType::Pointer VectorImage_to_imageOfVectors( typename VectorImageType::Pointer vectorImage )

{

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  typedef typename ImageOfVectorsType::Pointer ImageOfVectorsPointerType;

  typedef itk::Image<PixelType, Dimension> ImageType;

  // split channels

  std::vector<AntsImage<ImageType>> myChannels;

  AntsImage<VectorImageType> myCapsule = { vectorImage };

  myChannels = splitChannels<VectorImageType,ImageType>( myCapsule );

  // now merge channels into a vectorimage

  ImageOfVectorsPointerType finalImage = mergeChannels3<ImageType, ImageOfVectorsType>( myChannels );

  return finalImage;

}

template <typename RGBImageType, unsigned int Dimension>

AntsImage<itk::VectorImage<unsigned char, Dimension>> RgbToVector( AntsImage<RGBImageType> & antsImage )

{

  typedef typename RGBImageType::Pointer RGBImagePointerType;

  typedef itk::Image<itk::Vector<unsigned char,3>, Dimension>  VectorImageType;

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  RGBImagePointerType rgbImage = antsImage.ptr;//as< RGBImageType >( antsImage );

  // 2) Cast to Vector image for processing

  typedef itk::RGBToVectorImageAdaptor<RGBImageType> AdaptorInputType;

  typename AdaptorInputType::Pointer adaptInput = AdaptorInputType::New();

  adaptInput->SetImage( rgbImage );

  typedef itk::CastImageFilter<AdaptorInputType,VectorImageType> CastInputType;

  typename CastInputType::Pointer castInput = CastInputType::New();

  castInput->SetInput( adaptInput );

  castInput->Update();

  // get an image of vectors from rgb

  VectorImagePointerType vectorImage = castInput->GetOutput();

  // cast image of vectors to vectorimage

  typedef itk::VectorImage<unsigned char, Dimension>  VectorImageType2;

  typedef typename VectorImageType2::Pointer VectorImagePointerType2;

  VectorImagePointerType2 newVectorImage = imageOfVectors_to_VectorImage<VectorImageType2,VectorImageType,unsigned char, Dimension>(vectorImage);

  AntsImage<VectorImageType2> outImage = { newVectorImage };

  return outImage;

}

template <typename VectorImageType, unsigned int Dimension>

AntsImage<itk::Image< itk::RGBPixel<unsigned char>, Dimension >> VectorToRgb( AntsImage<VectorImageType> & antsImage )

{

  // cast VectorImage to ImageOfVectors

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  VectorImagePointerType vectorImage = antsImage.ptr;//as<VectorImageType>( antsImage );

  typedef itk::Image<itk::Vector<unsigned char,3>, Dimension>  VectorImageType2;

  typedef typename VectorImageType2::Pointer VectorImagePointerType2;

  VectorImagePointerType2 newVectorImage = VectorImage_to_imageOfVectors<VectorImageType,VectorImageType2,unsigned char, Dimension>(vectorImage);

  typedef itk::Image< itk::RGBPixel<unsigned char>, Dimension > RGBImageType;

  typedef typename RGBImageType::Pointer RGBImagePointerType;

  // 2) Cast to Vector image for processing

  typedef itk::VectorToRGBImageAdaptor<VectorImageType2> AdaptorInputType;

  typename AdaptorInputType::Pointer adaptInput = AdaptorInputType::New();

  adaptInput->SetImage( newVectorImage );

  typedef itk::CastImageFilter<AdaptorInputType,RGBImageType> CastInputType;

  typename CastInputType::Pointer castInput = CastInputType::New();

  castInput->SetInput( adaptInput );

  castInput->Update();

  //return wrap<RGBImageType>( castInput->GetOutput() );

  AntsImage<RGBImageType> outImage = { castInput->GetOutput() };

  return outImage;

}

void local_rgbToVector(nb::module_ &m)

{

  m.def("RgbToVector2", &RgbToVector<itk::Image<itk::RGBPixel<unsigned char>, 2>, 2>);

  m.def("RgbToVector3", &RgbToVector<itk::Image<itk::RGBPixel<unsigned char>, 3>, 3>);

  m.def("VectorToRgb2", &VectorToRgb<itk::VectorImage<unsigned char,2>, 2>);

  m.def("VectorToRgb3", &VectorToRgb<itk::VectorImage<unsigned char,3>, 3>);

}