#include <nanobind/nanobind.h>

#include <nanobind/stl/vector.h>

#include <nanobind/stl/string.h>

#include <nanobind/stl/tuple.h>

#include <nanobind/ndarray.h>

#include <exception>

#include <vector>

#include <string>

#include "itkImage.h"

#include "itkVectorImage.h"

#include "itkImageRegionIteratorWithIndex.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template< class ImageType, class VectorImageType >

AntsImage<VectorImageType> mergeChannels( std::vector<AntsImage<ImageType>> imageList )

{

  typedef typename ImageType::Pointer       ImagePointerType;

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  unsigned int nImages = imageList.size();

  std::vector<ImagePointerType> images;

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

  {

    images.push_back( imageList[i].ptr );

  }

  VectorImagePointerType vectorImage = VectorImageType::New();

  vectorImage->SetRegions( images[0]->GetLargestPossibleRegion() );

  vectorImage->SetSpacing( images[0]->GetSpacing() );

  vectorImage->SetOrigin( images[0]->GetOrigin() );

  vectorImage->SetDirection( images[0]->GetDirection() );

  vectorImage->SetNumberOfComponentsPerPixel( nImages );

  vectorImage->AllocateInitialized();

  // Fill image data

  itk::ImageRegionIteratorWithIndex<VectorImageType> it( vectorImage,

    vectorImage->GetLargestPossibleRegion() );

  while (!it.IsAtEnd() )

    {

    typename VectorImageType::PixelType pix;

    pix.SetSize( nImages );

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

      {

      pix[i] = images[i]->GetPixel(it.GetIndex());

      }

    vectorImage->SetPixel(it.GetIndex(), pix);

    ++it;

    }

    AntsImage<VectorImageType> outImage = { vectorImage };

  return outImage;

}

template< class ImageType, class VectorImageType >

typename VectorImageType::Pointer mergeChannels2( std::vector<AntsImage<ImageType>> imageList )

{

  typedef typename ImageType::Pointer       ImagePointerType;

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  unsigned int nImages = imageList.size();

  std::vector<ImagePointerType> images;

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

  {

    images.push_back( imageList[i].ptr );

  }

  VectorImagePointerType vectorImage = VectorImageType::New();

  vectorImage->SetRegions( images[0]->GetLargestPossibleRegion() );

  vectorImage->SetSpacing( images[0]->GetSpacing() );

  vectorImage->SetOrigin( images[0]->GetOrigin() );

  vectorImage->SetDirection( images[0]->GetDirection() );

  vectorImage->SetNumberOfComponentsPerPixel( nImages );

  vectorImage->AllocateInitialized();

  // Fill image data

  itk::ImageRegionIteratorWithIndex<VectorImageType> it( vectorImage,

    vectorImage->GetLargestPossibleRegion() );

  while (!it.IsAtEnd() )

    {

    typename VectorImageType::PixelType pix;

    pix.SetSize( nImages );

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

      {

      pix[i] = images[i]->GetPixel(it.GetIndex());

      }

    vectorImage->SetPixel(it.GetIndex(), pix);

    ++it;

    }

  //return wrap<VectorImageType>( vectorImage );

  //AntsImage<VectorImageType> outImage = { vectorImage };

  return vectorImage;

}

template< class ImageType, class VectorImageType >

typename VectorImageType::Pointer mergeChannels3( std::vector<AntsImage<ImageType>> imageList )

{

  typedef typename ImageType::Pointer       ImagePointerType;

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  unsigned int nImages = imageList.size();

  std::vector<ImagePointerType> images;

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

  {

    images.push_back( imageList[i].ptr );

  }

  VectorImagePointerType vectorImage = VectorImageType::New();

  vectorImage->SetRegions( images[0]->GetLargestPossibleRegion() );

  vectorImage->SetSpacing( images[0]->GetSpacing() );

  vectorImage->SetOrigin( images[0]->GetOrigin() );

  vectorImage->SetDirection( images[0]->GetDirection() );

  //vectorImage->SetNumberOfComponentsPerPixel( nImages );

  vectorImage->AllocateInitialized();

  // Fill image data

  itk::ImageRegionIteratorWithIndex<VectorImageType> it( vectorImage,

    vectorImage->GetLargestPossibleRegion() );

  while (!it.IsAtEnd() )

    {

    typename VectorImageType::PixelType pix;

    //pix.SetSize( nImages );

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

      {

      pix[i] = images[i]->GetPixel(it.GetIndex());

      }

    vectorImage->SetPixel(it.GetIndex(), pix);

    ++it;

    }

  //return wrap<VectorImageType>( vectorImage );

  //AntsImage<VectorImageType> outImage = { vectorImage };

  return vectorImage;

}

template< class VectorImageType, class ImageType>

std::vector<AntsImage<ImageType>> splitChannels( AntsImage<VectorImageType> & antsimage )

{

  typedef typename ImageType::Pointer       ImagePointerType;

  typedef typename VectorImageType::Pointer VectorImagePointerType;

  VectorImagePointerType input = antsimage.ptr;

  unsigned int nComponents = input->GetNumberOfComponentsPerPixel();

  // Create output images

  std::vector<ImagePointerType> images;

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

    {

    ImagePointerType image = ImageType::New();

    image->SetRegions( input->GetLargestPossibleRegion() );

    image->SetSpacing( input->GetSpacing() );

    image->SetOrigin( input->GetOrigin() );

    image->SetDirection( input->GetDirection() );

    image->AllocateInitialized();

    images.push_back( image );

    }

  // Fill image data

  itk::ImageRegionIteratorWithIndex<VectorImageType> it( input,

    input->GetLargestPossibleRegion() );

  while (!it.IsAtEnd() )

    {

    typename VectorImageType::PixelType pix = input->GetPixel(it.GetIndex());

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

      {

      images[i]->SetPixel(it.GetIndex(), pix[i]);

      }

    ++it;

    }

  std::vector<AntsImage<ImageType>> outputList( nComponents );

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

    {

        AntsImage<ImageType> tmpImage = { images[i] };

    outputList[i] = tmpImage;

    }

  return( outputList );

}

void local_mergeChannels(nb::module_ &m) {

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

}