#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 "itkCastImageFilter.h"
#include "antscore/itkSimulatedBSplineDisplacementFieldSource.h"
#include "antscore/itkSimulatedExponentialDisplacementFieldSource.h"
#include "antsImage.h"

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

template<class PrecisionType, unsigned int Dimension>
AntsImage<itk::VectorImage<PrecisionType, Dimension>> simulateBsplineDisplacementField(AntsImage<itk::Image<PrecisionType, Dimension>> & antsDomainImage,
                                             unsigned int numberOfRandomPoints,
                                             float standardDeviationDisplacementField,
                                             bool enforceStationaryBoundary,
                                             unsigned int numberOfFittingLevels,
                                             std::vector<unsigned int> numberOfControlPoints)
{
  using ImageType = itk::Image<PrecisionType, Dimension>;
  using ImagePointerType = typename ImageType::Pointer;

  ImagePointerType domainImage = antsDomainImage.ptr;

  using VectorType = itk::Vector<PrecisionType, Dimension>;
  using DisplacementFieldType = itk::Image<VectorType, Dimension>;
  using ANTsFieldType = itk::VectorImage<PrecisionType, Dimension>;
  using IteratorType = itk::ImageRegionIteratorWithIndex<DisplacementFieldType>;

  using BSplineSimulatorType = itk::SimulatedBSplineDisplacementFieldSource<DisplacementFieldType>;

  typename BSplineSimulatorType::ArrayType ncps;
  for( unsigned int d = 0; d < numberOfControlPoints.size(); ++d )
    {
    ncps = numberOfControlPoints[d];
    }

  using RealImageType = typename BSplineSimulatorType::RealImageType;
  using CastImageFilterType = itk::CastImageFilter<ImageType, RealImageType>;
  typename CastImageFilterType::Pointer caster = CastImageFilterType::New();
  caster->SetInput( domainImage );
  caster->Update();

  typename BSplineSimulatorType::Pointer bsplineSimulator = BSplineSimulatorType::New();
  bsplineSimulator->SetDisplacementFieldDomainFromImage( caster->GetOutput() );
  bsplineSimulator->SetNumberOfRandomPoints( numberOfRandomPoints );
  bsplineSimulator->SetEnforceStationaryBoundary( enforceStationaryBoundary );
  bsplineSimulator->SetDisplacementNoiseStandardDeviation( standardDeviationDisplacementField );
  bsplineSimulator->SetNumberOfFittingLevels( numberOfFittingLevels );
  bsplineSimulator->SetNumberOfControlPoints( ncps );
  bsplineSimulator->Update();

  typename ANTsFieldType::Pointer antsField = ANTsFieldType::New();
  antsField->CopyInformation( domainImage );
  antsField->SetRegions( domainImage->GetRequestedRegion() );
  antsField->SetVectorLength( Dimension );
  antsField->AllocateInitialized();

  IteratorType It( bsplineSimulator->GetOutput(),
    bsplineSimulator->GetOutput()->GetRequestedRegion() );
  for( It.GoToBegin(); !It.IsAtEnd(); ++It )
    {
    VectorType itkVector = It.Value();

    typename ANTsFieldType::PixelType antsVector( Dimension );
    for( unsigned int d = 0; d < Dimension; d++ )
      {
      antsVector[d] = itkVector[d];
      }
    antsField->SetPixel( It.GetIndex(), antsVector );
    }

  AntsImage<ANTsFieldType> outImage = { antsField };
  return outImage;
}

template<class PrecisionType, unsigned int Dimension>
AntsImage<itk::VectorImage<PrecisionType, Dimension>> simulateExponentialDisplacementField(AntsImage<itk::Image<PrecisionType, Dimension>> & antsDomainImage,
                                                 unsigned int numberOfRandomPoints,
                                                 float standardDeviationDisplacementField,
                                                 bool enforceStationaryBoundary,
                                                 float standardDeviationSmoothing)
{
  using ImageType = itk::Image<PrecisionType, Dimension>;
  using ImagePointerType = typename ImageType::Pointer;

  ImagePointerType domainImage = antsDomainImage.ptr;

  using VectorType = itk::Vector<PrecisionType, Dimension>;
  using DisplacementFieldType = itk::Image<VectorType, Dimension>;
  using ANTsFieldType = itk::VectorImage<PrecisionType, Dimension>;
  using IteratorType = itk::ImageRegionIteratorWithIndex<DisplacementFieldType>;

  using ExponentialSimulatorType = itk::SimulatedExponentialDisplacementFieldSource<DisplacementFieldType>;

  using RealImageType = typename ExponentialSimulatorType::RealImageType;
  using CastImageFilterType = itk::CastImageFilter<ImageType, RealImageType>;
  typename CastImageFilterType::Pointer caster = CastImageFilterType::New();
  caster->SetInput( domainImage );
  caster->Update();

  typename ExponentialSimulatorType::Pointer exponentialSimulator = ExponentialSimulatorType::New();
  exponentialSimulator->SetDisplacementFieldDomainFromImage( caster->GetOutput() );
  exponentialSimulator->SetNumberOfRandomPoints( numberOfRandomPoints );
  exponentialSimulator->SetEnforceStationaryBoundary( enforceStationaryBoundary );
  exponentialSimulator->SetDisplacementNoiseStandardDeviation( standardDeviationDisplacementField );
  exponentialSimulator->SetSmoothingStandardDeviation( standardDeviationSmoothing );
  exponentialSimulator->Update();

  typename ANTsFieldType::Pointer antsField = ANTsFieldType::New();
  antsField->CopyInformation( domainImage );
  antsField->SetRegions( domainImage->GetRequestedRegion() );
  antsField->SetVectorLength( Dimension );
  antsField->AllocateInitialized();

  IteratorType It( exponentialSimulator->GetOutput(),
    exponentialSimulator->GetOutput()->GetRequestedRegion() );
  for( It.GoToBegin(); !It.IsAtEnd(); ++It )
    {
    VectorType itkVector = It.Value();

    typename ANTsFieldType::PixelType antsVector( Dimension );
    for( unsigned int d = 0; d < Dimension; d++ )
      {
      antsVector[d] = itkVector[d];
      }
    antsField->SetPixel( It.GetIndex(), antsVector );
    }

  AntsImage<ANTsFieldType> outImage = { antsField };
  return outImage;
}

void local_simulateDisplacementField(nb::module_ &m)
{
  m.def("simulateBsplineDisplacementField2D", &simulateBsplineDisplacementField<float, 2>);
  m.def("simulateBsplineDisplacementField3D", &simulateBsplineDisplacementField<float, 3>);

  m.def("simulateExponentialDisplacementField2D", &simulateExponentialDisplacementField<float, 2>);
  m.def("simulateExponentialDisplacementField3D", &simulateExponentialDisplacementField<float, 3>);
}