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

#include "itkImageRegionIteratorWithIndex.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template< unsigned int Dimension >

nb::dict labelStatsHelper(

  typename itk::Image< float, Dimension >::Pointer image,

  typename itk::Image< unsigned int, Dimension>::Pointer labelImage)

{

  typedef float PixelType;

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

  typedef unsigned int LabelType;

  typedef typename ImageType::PointType PointType;

  typedef itk::Image< LabelType, Dimension > LabelImageType;

  typedef itk::ImageRegionIteratorWithIndex<LabelImageType>                    Iterator;

  typedef itk::LabelStatisticsImageFilter< ImageType, LabelImageType >

    LabelStatisticsImageFilterType;

  typename LabelStatisticsImageFilterType::Pointer labelStatisticsImageFilter =

    LabelStatisticsImageFilterType::New();

  labelStatisticsImageFilter->SetInput( image );

  labelStatisticsImageFilter->SetLabelInput( labelImage );

  labelStatisticsImageFilter->Update();

  typedef typename LabelStatisticsImageFilterType::ValidLabelValuesContainerType

    ValidLabelValuesType;

  long nlabs = labelStatisticsImageFilter->GetNumberOfLabels();

  std::vector<double> labelvals(nlabs);

  std::vector<double> means(nlabs);

  std::vector<double> mins(nlabs);

  std::vector<double> maxes(nlabs);

  std::vector<double> variances(nlabs);

  std::vector<double> counts(nlabs);

  std::vector<double> volumes(nlabs);

  std::vector<double> x(nlabs);

  std::vector<double> y(nlabs);

  std::vector<double> z(nlabs);

  std::vector<double> t(nlabs);

  std::vector<double> mass(nlabs,0.0);

  typename ImageType::SpacingType spacing = image->GetSpacing();

  float voxelVolume = 1.0;

  for (unsigned int ii = 0; ii < spacing.Size(); ii++)

  {

    voxelVolume *= spacing[ii];

  }

  std::map<LabelType, LabelType> RoiList;

  LabelType ii = 0; // counter for label values

  for (typename ValidLabelValuesType::const_iterator

         labelIterator  = labelStatisticsImageFilter->GetValidLabelValues().begin();

         labelIterator != labelStatisticsImageFilter->GetValidLabelValues().end();

         ++labelIterator)

  {

    if ( labelStatisticsImageFilter->HasLabel(*labelIterator) )

    {

      LabelType labelValue = *labelIterator;

      labelvals[ii] = labelValue;

      means[ii]     = labelStatisticsImageFilter->GetMean(labelValue);

      mins[ii]      = labelStatisticsImageFilter->GetMinimum(labelValue);

      maxes[ii]     = labelStatisticsImageFilter->GetMaximum(labelValue);

      variances[ii] = labelStatisticsImageFilter->GetVariance(labelValue);

      counts[ii]    = labelStatisticsImageFilter->GetCount(labelValue);

      volumes[ii]   = labelStatisticsImageFilter->GetCount(labelValue) * voxelVolume;

      RoiList[ labelValue ] = ii;

    }

    ++ii;

  }

  Iterator It( labelImage, labelImage->GetLargestPossibleRegion() );

  std::vector<PointType> comvec;

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

    {

    typename ImageType::PointType myCenterOfMass;

    myCenterOfMass.Fill(0);

    comvec.push_back( myCenterOfMass );

    }

  for( It.GoToBegin(); !It.IsAtEnd(); ++It )

    {

    LabelType label = static_cast<LabelType>( It.Get() );

    if(  label > 0  )

      {

      typename ImageType::PointType point;

      image->TransformIndexToPhysicalPoint( It.GetIndex(), point );

      for( unsigned int i = 0; i < spacing.Size(); i++ )

        {

        comvec[  RoiList[ label ] ][i] += point[i];

        }

      mass[  RoiList[ label ] ] += image->GetPixel( It.GetIndex() );

      }

    }

  for ( unsigned int labelcount = 0; labelcount < comvec.size(); labelcount++ )

    {

    for ( unsigned int k = 0; k < Dimension; k++ )

      {

      comvec[ labelcount ][k] = comvec[ labelcount ][k] / counts[labelcount];

      }

    x[labelcount]=comvec[ labelcount ][0];

    y[labelcount]=comvec[ labelcount ][1];

    if ( Dimension > 2 ) z[labelcount]=comvec[ labelcount ][2];

    if ( Dimension > 3 ) t[labelcount]=comvec[ labelcount ][3];

    }

  nb::dict labelStats;

  labelStats["LabelValue"] = labelvals;

  labelStats["Mean"] = means;

  labelStats["Min"] = mins;

  labelStats["Max"] = maxes;

  labelStats["Variance"] = variances;

  labelStats["Count"] = counts;

  labelStats["Volume"] = volumes;

  labelStats["Mass"] = mass;

  labelStats["x"] = x;

  labelStats["y"] = y;

  labelStats["z"] = z;

  labelStats["t"] = t;

  return (labelStats);

}

template <unsigned int Dimension>

nb::dict labelStats(AntsImage<itk::Image<float, Dimension>> & py_image,

                    AntsImage<itk::Image<unsigned int, Dimension>> & py_labelImage)

{ 

  typedef itk::Image<float, Dimension> FloatImageType;

  typedef itk::Image<unsigned int, Dimension> IntImageType;

  typedef typename FloatImageType::Pointer FloatImagePointerType;

  typedef typename IntImageType::Pointer IntImagePointerType;

  FloatImagePointerType myimage = py_image.ptr;

  IntImagePointerType mylabelimage = py_labelImage.ptr;

  return labelStatsHelper<Dimension>( myimage, mylabelimage );

}

void local_labelStats(nb::module_ &m)

{

  m.def("labelStats2D", &labelStats<2>);

  m.def("labelStats3D", &labelStats<3>);

  m.def("labelStats4D", &labelStats<4>);

}