#include <nanobind/nanobind.h>

#include <nanobind/stl/vector.h>

#include <nanobind/stl/string.h>

#include <nanobind/stl/tuple.h>

#include <nanobind/ndarray.h>

#include <tuple>

#include "itkImageFileWriter.h"

#include "itkImageIOBase.h"

#include "itkNiftiImageIOFactory.h"

#include "itkMetaImageIOFactory.h"

#include "itkImageFileReader.h"

#include "itkImage.h"

#include "itkPyBuffer.h"

#include "readImage.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template <typename ImageType>

AntsImage<ImageType> imageRead( std::string filename )

{

    typedef typename ImageType::Pointer           ImagePointerType;

    typedef itk::ImageFileReader< ImageType >     ImageReaderType;

    typename ImageReaderType::Pointer image_reader = ImageReaderType::New() ;

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

    image_reader->Update();

    ImagePointerType itkImage = image_reader->GetOutput();

    AntsImage<ImageType> myImage = { itkImage };

    return myImage;

}

/**

 * Create an ANTsImage from a numpy array

 *

 * The data array must be C-contiguous, and ordered correctly

 * See "4.1.7 Importing Image Data from a Buffer"

 * https://itk.org/ITKSoftwareGuide/html/Book1/ITKSoftwareGuide-Book1ch4.html

 */

template <typename ImageType>

AntsImage<ImageType> fromNumpy( nb::ndarray<nb::numpy> data, nb::tuple datashape )

{

    typedef typename ImageType::Pointer ImagePointerType;

    ImagePointerType antsImage = ImageType::New();

    typedef itk::PyBuffer<ImageType> PyBufferType;

    nb::object o = nb::cast(data);

    antsImage = PyBufferType::_GetImageViewFromArray(o.ptr(), datashape.ptr(), nb::make_tuple(1)[0].ptr());

    AntsImage<ImageType> myImage = { antsImage };

    return myImage;

}

void local_readImage(nb::module_ &m) {

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

}