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

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template < typename ImageType >

AntsImage<ImageType> padImage( AntsImage<ImageType> & antsImage,

                      std::vector<int> lowerPadDims,

                      std::vector<int> upperPadDims,

                      float padValue )

{

  typedef typename ImageType::Pointer ImagePointerType;

  typename ImageType::Pointer itkImage = antsImage.ptr;

  typename ImageType::SizeType lowerExtendRegion;

  lowerExtendRegion[0] = lowerPadDims[0];

  lowerExtendRegion[1] = lowerPadDims[1];

  if (lowerPadDims.size() == 3)

  {

    lowerExtendRegion[2] = lowerPadDims[2];

  }

  typename ImageType::SizeType upperExtendRegion;

  upperExtendRegion[0] = upperPadDims[0];

  upperExtendRegion[1] = upperPadDims[1];

  if (upperPadDims.size() == 3)

  {

    upperExtendRegion[2] = upperPadDims[2];

  }

  //ImageType::PixelType constantPixel = padValue;

  typedef itk::ConstantPadImageFilter<ImageType, ImageType> PadImageFilterType;

  typename PadImageFilterType::Pointer padFilter = PadImageFilterType::New();

  padFilter->SetInput( itkImage );

  padFilter->SetPadLowerBound( lowerExtendRegion );

  padFilter->SetPadUpperBound( upperExtendRegion );

  padFilter->SetConstant( padValue );

  padFilter->Update();

  FixNonZeroIndex<ImageType>( padFilter->GetOutput() );

  AntsImage<ImageType> myImage = { padFilter->GetOutput() };

  return myImage;

}

void local_padImage(nb::module_ &m) 

{

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

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

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

}