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

#include "itkSaltAndPepperNoiseImageFilter.h"

#include "itkShotNoiseImageFilter.h"

#include "itkSpeckleNoiseImageFilter.h"

#include "antsImage.h"

namespace nb = nanobind;

using namespace nb::literals;

template <typename ImageType>

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

                                   float mean,

                                   float standardDeviation )

{

  typename ImageType::Pointer itkImage = antsImage.ptr;

  using NoiseFilterType = itk::AdditiveGaussianNoiseImageFilter<ImageType, ImageType>;

  typename NoiseFilterType::Pointer noiser = NoiseFilterType::New();

  noiser->SetInput( itkImage );

  noiser->SetMean( mean );

  noiser->SetStandardDeviation( standardDeviation );

  noiser->Update();

  AntsImage<ImageType> outImage = { noiser->GetOutput() };

  return outImage;

}

template <typename ImageType>

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

                                float probability,

                                float saltValue,

                                float pepperValue )

{

  typename ImageType::Pointer itkImage = antsImage.ptr;

  using NoiseFilterType = itk::SaltAndPepperNoiseImageFilter<ImageType, ImageType>;

  typename NoiseFilterType::Pointer noiser = NoiseFilterType::New();

  noiser->SetInput( itkImage );

  noiser->SetProbability( probability );

  noiser->SetSaltValue( saltValue );

  noiser->SetPepperValue( pepperValue );

  noiser->Update();

  AntsImage<ImageType> outImage = { noiser->GetOutput() };

  return outImage;

}

template <typename ImageType>

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

                       float scale

 )

{

  typename ImageType::Pointer itkImage = antsImage.ptr;

  using NoiseFilterType = itk::ShotNoiseImageFilter<ImageType, ImageType>;

  typename NoiseFilterType::Pointer noiser = NoiseFilterType::New();

  noiser->SetInput( itkImage );

  noiser->SetScale( scale );

  noiser->Update();

  AntsImage<ImageType> outImage = { noiser->GetOutput() };

  return outImage;

}

template <typename ImageType>

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

                          float scale

 )

{

  typename ImageType::Pointer itkImage = antsImage.ptr;

  using NoiseFilterType = itk::SpeckleNoiseImageFilter<ImageType, ImageType>;

  typename NoiseFilterType::Pointer noiser = NoiseFilterType::New();

  noiser->SetInput( itkImage );

  noiser->SetStandardDeviation( scale );

  noiser->Update();

  AntsImage<ImageType> outImage = { noiser->GetOutput() };

  return outImage;

}

void local_addNoiseToImage(nb::module_ &m)

{

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

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

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

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

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

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

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

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

}