/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkRealFFTTest_h
#define itkRealFFTTest_h

/* This test is built for filters specialized for real-to-complex
 * forward and complex-to-real inverse Fast Fourier Transforms using
 * VNL and FFTW FFT libraries. */
#include "itkConfigure.h"
#include "itkImage.h"
#include "itkImageRegionIterator.h"
#include "itkVnlRealToHalfHermitianForwardFFTImageFilter.h"
#include "itkVnlHalfHermitianToRealInverseFFTImageFilter.h"

#if defined(ITK_USE_FFTWF) || defined(ITK_USE_FFTWD)
#  include "itkFFTWRealToHalfHermitianForwardFFTImageFilter.h"
#  include "itkFFTWHalfHermitianToRealInverseFFTImageFilter.h"
#endif

#include "itksys/SystemTools.hxx"
#include "vnl/vnl_sample.h"
#include <cmath>

/* test_fft is the test function and it is templated over the pixel, Image
 * dimensions and the FFT library to be used. */
template <typename TPixel, unsigned int VImageDimensions, typename R2CType, typename C2RType>
int
test_fft(unsigned int * SizeOfDimensions)
{
  using RealImageType = itk::Image<TPixel, VImageDimensions>;
  using ComplexImageType = itk::Image<std::complex<TPixel>, VImageDimensions>;
  unsigned int                      counter = 0;
  typename RealImageType::SizeType  imageSize;
  typename RealImageType::IndexType imageIndex;

  // We are testing the FFT for 1D, 2D, and 3D images. An array
  // (SizeOfDimensions) containing the sizes of each dimension is
  // passed as an argument to this function. Based on the template
  // argument VImageDimensions, we create a 1D, 2D, or 3D image by
  // selecting the sizes of image dimensions from this array.
  for (unsigned int i = 0; i < VImageDimensions; ++i)
  {
    imageSize.SetElement(i, SizeOfDimensions[i]);
    imageIndex.SetElement(i, -2 * static_cast<int>(i) - 5); // Test for handling non-zero
                                                            // image indices correctly
  }

  typename RealImageType::RegionType region{ imageIndex, imageSize };

  // Create the Real Image.
  auto realImage = RealImageType::New();
  realImage->SetRegions(region);
  realImage->Allocate();
  vnl_sample_reseed(123456);

  // We use 2 region iterators for this test: the original image
  // iterator and another iterator for the resultant image after
  // performing FFT and IFFT.
  itk::ImageRegionIterator<RealImageType> originalImageIterator(realImage, region);

  // Allocate random pixel values to the image by iterating through it
  // and print out the image data.
  try
  {
    std::cerr << "---- Original image ----" << std::endl;
    while (!originalImageIterator.IsAtEnd())
    {
      TPixel val = vnl_sample_uniform(0.0, 16384.0);
      if ((counter + 1) % SizeOfDimensions[0] == 0)
      {
        std::cout << val << std::endl;
      }
      else
      {
        std::cout << val << ' ';
      }
      counter++;
      originalImageIterator.Set(val);
      ++originalImageIterator;
    }
    std::cout << std::endl << std::endl;
  }
  catch (const itk::ExceptionObject & ex)
  {
    ex.Print(std::cerr);
    return -1;
  }

  // Real to complex pointer. This computes the forward FFT.
  auto R2C = R2CType::New();

  // Complex to Real pointer. This computes the HalfHermitianToRealInverse FFT.
  auto C2R = C2RType::New();

  // Set the real image created as the input to the forward FFT
  // filter.
  R2C->SetInput(realImage);
  R2C->Print(std::cout);

  try
  {
    R2C->Update();
  }
  catch (const itk::ExceptionObject & ex)
  {
    ex.Print(std::cerr);
    return -1;
  }

  // Get the size and the pointer to the complex image.
  const typename ComplexImageType::Pointer    complexImage = R2C->GetOutput();
  std::complex<TPixel> *                      fftbuf = complexImage->GetBufferPointer();
  const typename ComplexImageType::SizeType & complexImageSize = complexImage->GetLargestPossibleRegion().GetSize();

  unsigned int sizes[4] = { 1, 1, 1, 1 };
  for (unsigned int i = 0; i < VImageDimensions; ++i)
  {
    sizes[i] = complexImageSize[i];
  }

  /* Print out the frequency domain data obtained after performing
   * the forward transform. */
  std::cout << "Frequency domain data after forward transform:" << std::endl;
  for (unsigned int i = 0; i < sizes[2]; ++i)
  {
    const unsigned int zStride = i * sizes[1] * sizes[0];
    for (unsigned int j = 0; j < sizes[1]; ++j)
    {
      const unsigned int yStride = j * sizes[0];
      for (unsigned int k = 0; k < sizes[0]; ++k)
      {
        std::cout << fftbuf[zStride + yStride + k] << ' ';
      }
      std::cout << std::endl;
    }
  }

  std::cout << std::endl << std::endl;

  // Perform the HalfHermitianToRealInverse FFT to get back the Real Image. C2R is the
  // complex conjugate to real image filter and we give the resulting
  // complex image as input to this filter. This is the HalfHermitianToRealInverse FFT of
  // the image.
  C2R->SetInput(complexImage);

  // Inform the filter that there's an odd # of pixels in the x
  // dimension.
  C2R->SetActualXDimensionIsOddInput(R2C->GetActualXDimensionIsOddOutput());
  C2R->Print(std::cout);
  C2R->Update();
  std::cerr << "C2R region: " << C2R->GetOutput()->GetLargestPossibleRegion() << std::endl;
  const typename RealImageType::Pointer imageAfterHalfHermitianToRealInverseFFT = C2R->GetOutput();

  // The HalfHermitianToRealInverse FFT image iterator is the resultant iterator after we
  // perform the FFT and HalfHermitianToRealInverse FFT on the Original Image. */
  itk::ImageRegionIterator<RealImageType> inverseFFTImageIterator(imageAfterHalfHermitianToRealInverseFFT, region);
  counter = 0;
  inverseFFTImageIterator.GoToBegin();

  // Print the Image data obtained by performing the HalfHermitianToRealInverse FFT.
  std::cerr << "---- HalfHermitianToRealInverse FFT image ----" << std::endl;
  while (!inverseFFTImageIterator.IsAtEnd())
  {
    TPixel val = inverseFFTImageIterator.Value();
    if ((counter + 1) % SizeOfDimensions[0] == 0)
    {
      std::cerr << val << std::endl;
    }
    else
    {
      std::cerr << val << ' ';
    }
    counter++;
    ++inverseFFTImageIterator;
  }

  std::cerr << std::endl << std::endl;

  // Subtract the Original image Pixel Values from the resultant image
  // values and test whether they are greater than 0.01 for the test
  // to pass.
  originalImageIterator.GoToBegin();
  inverseFFTImageIterator.GoToBegin();
  while (!originalImageIterator.IsAtEnd())
  {
    TPixel val = originalImageIterator.Value();
    TPixel val2 = inverseFFTImageIterator.Value();
    TPixel diff = itk::Math::Absolute(val - val2);
    if (itk::Math::NotAlmostEquals(val, TPixel{}))
    {
      diff /= itk::Math::Absolute(val);
    }
    if (diff > 0.01)
    {
      std::cerr << "Diff found in test_fft: " << val << ' ' << val2 << " diff " << diff << std::endl;
      return -1;
    }
    ++originalImageIterator;
    ++inverseFFTImageIterator;
  }
  std::cout << std::endl << std::endl;

  return 0;
}


/* test_fft_rtc is the test function to compare two implementations
 * (Direct FFT only).  It is templated over the pixel, Image
 * dimensions and the FFT libraries to be used. */
template <typename TPixel, unsigned int VImageDimensions, typename R2CAType, typename R2CBType>
int
test_fft_rtc(unsigned int * SizeOfDimensions)
{
  using RealImageType = itk::Image<TPixel, VImageDimensions>;
  using ComplexImageType = itk::Image<std::complex<TPixel>, VImageDimensions>;
  unsigned int                      counter = 0;
  typename RealImageType::SizeType  imageSize;
  typename RealImageType::IndexType imageIndex;

  // We are testing the FFT for 1D, 2D, and 3D images. An array
  // (SizeOfDimensions) containing the sizes of each dimension is
  // passed as an argument to this function. Based on the template
  // argument VImageDimensions, we create a 1D, 2D, or 3D image by
  // selecting the sizes of image dimensions from this array.
  for (unsigned int i = 0; i < VImageDimensions; ++i)
  {
    imageSize.SetElement(i, SizeOfDimensions[i]);
    imageIndex.SetElement(i, 0);
  }

  typename RealImageType::RegionType region{ imageIndex, imageSize };
  auto                               realImage = RealImageType::New();

  // Create the Real Image.
  realImage->SetRegions(region);
  realImage->Allocate();
  vnl_sample_reseed(static_cast<int>(itksys::SystemTools::GetTime() / 10000.0));

  // We use 2 region iterators for this test the original image
  // iterator and another iterator for the resultant image after
  // performing FFT and IFFT.
  itk::ImageRegionIterator<RealImageType> originalImageIterator(realImage, region);

  // Allocate random pixel values to the image by iterating through it
  // and Print out the image data.
  try
  {
    while (!originalImageIterator.IsAtEnd())
    {
      TPixel val = vnl_sample_uniform(0.0, 16384.0);
      if ((counter + 1) % SizeOfDimensions[0] == 0)
      {
        std::cout << val << std::endl;
      }
      else
      {
        std::cout << val << ' ';
      }
      counter++;
      originalImageIterator.Set(val);
      ++originalImageIterator;
    }
    std::cout << std::endl << std::endl;
  }
  catch (const itk::ExceptionObject & ex)
  {
    ex.Print(std::cerr);
    return -1;
  }

  // Real to complex pointers. This computes the forward FFT.
  auto R2Ca = R2CAType::New();

  // Real to complex pointers. This computes the forward FFT.
  auto R2Cb = R2CBType::New();

  // Set the real image created as the input to the forward FFT
  // filter.
  R2Ca->SetInput(realImage);
  R2Ca->Update();

  R2Cb->SetInput(realImage);
  R2Cb->Update();

  // Get the size and the pointer to the complex image.
  const typename ComplexImageType::Pointer    complexImageA = R2Ca->GetOutput();
  std::complex<TPixel> *                      fftbufA = complexImageA->GetBufferPointer();
  const typename ComplexImageType::SizeType & complexImageSizeA = complexImageA->GetLargestPossibleRegion().GetSize();

  const typename ComplexImageType::Pointer    complexImageB = R2Cb->GetOutput();
  std::complex<TPixel> *                      fftbufB = complexImageB->GetBufferPointer();
  const typename ComplexImageType::SizeType & complexImageSizeB = complexImageB->GetLargestPossibleRegion().GetSize();


  unsigned int sizesA[4] = { 1, 1, 1, 1 };
  unsigned int sizesB[4] = { 1, 1, 1, 1 };
  for (unsigned int i = 0; i < VImageDimensions; ++i)
  {
    // The size may be different if one implementation returns a
    // full matrix but not the other.
    sizesA[i] = complexImageSizeA[i];
    sizesB[i] = complexImageSizeB[i];
  }

  // Print out the frequency domain data obtained after performing
  // the forward transform.
  std::cout << "Frequency domain data after forward transform:" << std::endl;
  for (unsigned int i = 0; i < sizesA[2]; ++i)
  {
    const unsigned int zStride = i * sizesA[1] * sizesA[0];
    for (unsigned int j = 0; j < sizesA[1]; ++j)
    {
      const unsigned int yStride = j * sizesA[0];
      for (unsigned int k = 0; k < sizesA[0]; ++k)
      {
        std::cout << fftbufA[zStride + yStride + k] << ' ';
      }
      std::cout << std::endl;
    }
  }
  std::cout << std::endl << std::endl;

  for (unsigned int i = 0; i < sizesB[2]; ++i)
  {
    const unsigned int zStride = i * sizesB[1] * sizesB[0];
    for (unsigned int j = 0; j < sizesB[1]; ++j)
    {
      const unsigned int yStride = j * sizesB[0];
      for (unsigned int k = 0; k < sizesB[0]; ++k)
      {
        std::cout << fftbufB[zStride + yStride + k] << ' ';
      }
      std::cout << std::endl;
    }
  }
  std::cout << std::endl << std::endl;


  // Subtract the pixel values from the two images. If one pixel
  // difference is greater than 0.01, the test is considered to have
  // failed.
  for (unsigned int i = 0; i < std::min(sizesA[2], sizesB[2]); ++i)
  {
    const unsigned int zStrideA = i * sizesA[1] * sizesA[0];
    const unsigned int zStrideB = i * sizesB[1] * sizesB[0];
    for (unsigned int j = 0; j < std::min(sizesA[1], sizesB[1]); ++j)
    {
      const unsigned int yStrideA = j * sizesA[0];
      const unsigned int yStrideB = j * sizesB[0];
      for (unsigned int k = 0; k < std::min(sizesA[0], sizesB[0]); ++k)
      {
        const double val = itk::Math::Absolute(fftbufA[zStrideA + yStrideA + k]);
        double       diff = itk::Math::Absolute(fftbufA[zStrideA + yStrideA + k] - fftbufB[zStrideB + yStrideB + k]);
        if (itk::Math::NotAlmostEquals(val, 0.0))
        {
          diff /= itk::Math::Absolute(val);
        }
        if (diff > 0.01)
        {
          std::cerr << "Diff found in test_fft_r2c: " << fftbufA[zStrideA + yStrideA + k] << ' '
                    << fftbufB[zStrideB + yStrideB + k] << " diff " << diff << std::endl;
          return -1;
        }
      }
    }
  }

  std::cout << std::endl << std::endl;

  return 0;
}
#endif
