#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 <exception>

#include <vector>

#include <string>

#include "itkImage.h"

#include "itkLabelOverlapMeasuresImageFilter.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template<class PrecisionType, unsigned int ImageDimension>

nb::dict labelOverlapMeasures( AntsImage<itk::Image<PrecisionType, ImageDimension>> &  antsSourceImage,

                                AntsImage<itk::Image<PrecisionType, ImageDimension>> &  antsTargetImage )

{

  using ImageType = itk::Image<PrecisionType, ImageDimension>;

  using ImagePointerType = typename ImageType::Pointer;

  typename ImageType::Pointer itkSourceImage = antsSourceImage.ptr;

  typename ImageType::Pointer itkTargetImage = antsTargetImage.ptr;

  using FilterType = itk::LabelOverlapMeasuresImageFilter<ImageType>;

  typename FilterType::Pointer filter = FilterType::New();

  filter->SetSourceImage( itkSourceImage );

  filter->SetTargetImage( itkTargetImage );

  filter->Update();

  typename FilterType::MapType labelMap = filter->GetLabelSetMeasures();

  // Sort the labels

  std::vector<PrecisionType> allLabels;

  allLabels.clear();

  for( typename FilterType::MapType::const_iterator it = labelMap.begin();

       it != labelMap.end(); ++it )

    {

    if( (*it).first == 0 )

      {

      continue;

      }

    const int label = (*it).first;

    allLabels.push_back( label );

    }

  std::sort( allLabels.begin(), allLabels.end() );

  // Now put the results in an Rcpp data frame

  unsigned int vectorLength = 1 + allLabels.size();

  std::vector<PrecisionType> labels( vectorLength );

  std::vector<double> totalOrTargetOverlap( vectorLength );

  std::vector<double> unionOverlap( vectorLength );

  std::vector<double> meanOverlap( vectorLength );

  std::vector<double> volumeSimilarity( vectorLength );

  std::vector<double> falseNegativeError( vectorLength );

  std::vector<double> falsePositiveError( vectorLength );

  // We'll replace label '0' with "All" in the R wrapper.

  labels[0] = itk::NumericTraits<PrecisionType>::Zero;

  totalOrTargetOverlap[0] = filter->GetTotalOverlap();

  unionOverlap[0] = filter->GetUnionOverlap();

  meanOverlap[0] = filter->GetMeanOverlap();

  volumeSimilarity[0] = filter->GetVolumeSimilarity();

  falseNegativeError[0] = filter->GetFalseNegativeError();

  falsePositiveError[0] = filter->GetFalsePositiveError();

  unsigned int i = 1;

  typename std::vector<PrecisionType>::const_iterator itL = allLabels.begin();

  for( itL = allLabels.begin(); itL != allLabels.end(); ++itL )

    {

    labels[i] = *itL;

    totalOrTargetOverlap[i] = filter->GetTargetOverlap( *itL );

    unionOverlap[i] = filter->GetUnionOverlap( *itL );

    meanOverlap[i] = filter->GetMeanOverlap( *itL );

    volumeSimilarity[i] = filter->GetVolumeSimilarity( *itL );

    falseNegativeError[i] = filter->GetFalseNegativeError( *itL );

    falsePositiveError[i] = filter->GetFalsePositiveError( *itL );

    i++;

    }

  nb::dict labelOverlapMeasures;

  labelOverlapMeasures["Label"] = labels;

  labelOverlapMeasures["TotalOrTargetOverlap"] = totalOrTargetOverlap;

  labelOverlapMeasures["UnionOverlap"] = unionOverlap;

  labelOverlapMeasures["MeanOverlap"] = meanOverlap;

  labelOverlapMeasures["VolumeSimilarity"] = volumeSimilarity;

  labelOverlapMeasures["FalseNegativeError"] = falseNegativeError;

  labelOverlapMeasures["FalsePositiveError"] = falsePositiveError;

  return labelOverlapMeasures;

}

void local_labelOverlapMeasures(nb::module_ &m)

{

  m.def("labelOverlapMeasures2D", &labelOverlapMeasures<unsigned int, 2>);

  m.def("labelOverlapMeasures3D", &labelOverlapMeasures<unsigned int, 3>);

  m.def("labelOverlapMeasures4D", &labelOverlapMeasures<unsigned int, 4>);

}