#include <nanobind/nanobind.h>
#include <nanobind/stl/vector.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/tuple.h>
#include <nanobind/ndarray.h>

#include "itkImage.h"
#include "itkExtractImageFilter.h"

#include "antsImage.h"

namespace nb = nanobind;
using namespace nb::literals;

template < typename ImageType, typename SliceImageType>
AntsImage<SliceImageType> sliceImage( AntsImage<ImageType> & antsImage, int plane, int slice, int collapseStrategy)
{
    typename ImageType::Pointer itkImage = antsImage.ptr;

    typedef itk::ExtractImageFilter< ImageType, SliceImageType > FilterType;
    typename FilterType::Pointer filter = FilterType::New();

    typename ImageType::RegionType inputRegion = itkImage->GetLargestPossibleRegion();
    typename ImageType::SizeType size = inputRegion.GetSize();
    size[plane] = 0;

    typename ImageType::IndexType start = inputRegion.GetIndex();
    const unsigned int sliceNumber = slice;
    start[plane] = sliceNumber;

    typename ImageType::RegionType desiredRegion;
    desiredRegion.SetSize( size );
    desiredRegion.SetIndex( start );

    filter->SetExtractionRegion( desiredRegion );
    filter->SetInput( itkImage );
    if( collapseStrategy == 0 )
      { 
      filter->SetDirectionCollapseToSubmatrix();
      } 
    else if( collapseStrategy == 1 )
      {
      filter->SetDirectionCollapseToIdentity();
      } 
    else // if( collapseStrategy == 2 ) 
      {
      filter->SetDirectionCollapseToGuess();
      } 
  
    filter->Update();

    AntsImage<SliceImageType> myImage = { filter->GetOutput() };
    return myImage;

}

void local_sliceImage(nb::module_ &m) 
{
    m.def("sliceImage", &sliceImage<itk::Image<float,3>, itk::Image<float,2>>);
    m.def("sliceImage", &sliceImage<itk::Image<float,4>, itk::Image<float,3>>);
}