/*=========================================================================
 *
 *  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
 *
 *         http://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 itkDirectFourierReconstructionImageToImageFilter_hxx
#define itkDirectFourierReconstructionImageToImageFilter_hxx

#include "itkDirectFourierReconstructionImageToImageFilter.h"

namespace itk
{
/**
 * Initialize member variables with meaningful values.
 */
template <typename TInputImage, typename TOutputImage>
DirectFourierReconstructionImageToImageFilter<TInputImage,
                                              TOutputImage>::DirectFourierReconstructionImageToImageFilter()
  : Superclass()
{
  constexpr double RADIANS = 1.0;

  m_ZeroPadding = 2;
  m_OverSampling = 2;
  m_Cutoff = 1.0;
  m_AlphaRange = 180;

  m_ZDirection = 1;
  m_RDirection = 0;
  m_AlphaDirection = 2;

  m_RadialSplineOrder = 3;

  m_PI = 4 * std::atan(RADIANS);
}

/**
 * Print out class state (member variables)
 */
template <typename TInputImage, typename TOutputImage>
void
DirectFourierReconstructionImageToImageFilter<TInputImage, TOutputImage>::PrintSelf(std::ostream & os,
                                                                                    Indent         indent) const
{
  // call the superclass' implementation of this method
  Superclass::PrintSelf(os, indent);

  os << indent << "Zero Padding Factor: " << this->GetZeroPadding() << std::endl;
  os << indent << "Fourier Oversampling Factor: " << this->GetOverSampling() << std::endl;
  os << indent << "Radial Spline Order: " << this->GetRadialSplineOrder() << std::endl;
  os << indent << "Fourier Radial Cutoff Frequency: " << this->GetCutoff() << std::endl;
  os << indent << "Alpha Range: " << this->GetAlphaRange() << std::endl;
  os << indent << "Z Direction: " << this->GetZDirection() << std::endl;
  os << indent << "Alpha Direction: " << this->GetAlphaDirection() << std::endl;
  os << indent << "Radial Direction: " << this->GetRDirection() << std::endl;
  os << indent << "Input Requested Region: " << m_InputRequestedRegion << std::endl;
}

/**
 * Calculate image boundaries and define output regions, spacing, origin etc.
 */
template <typename TInputImage, typename TOutputImage>
void
DirectFourierReconstructionImageToImageFilter<TInputImage, TOutputImage>::GenerateOutputInformation()
{
  // call the superclass' implementation of this method
  Superclass::GenerateOutputInformation();

  ConstInputImagePointer inputImage = this->GetInput();
  OutputImagePointer     outputImage = this->GetOutput();

  if (!inputImage || !outputImage)
  {
    return;
  }

  RegionType inputRegion = inputImage->GetLargestPossibleRegion();
  IndexType  inputStart = inputRegion.GetIndex();
  SizeType   inputSize = inputRegion.GetSize();

  IndexType outputStart;
  SizeType  outputSize;

  // Conserve axial dimensions
  outputSize[2] = inputSize[m_ZDirection];
  outputStart[2] = inputStart[m_ZDirection];

  // Radial dimensions give in-plane extent
  outputSize[0] = outputSize[1] = inputSize[m_RDirection];
  outputStart[0] = outputStart[1] = 0;

  RegionType outputRegion;
  outputRegion.SetIndex(outputStart);
  outputRegion.SetSize(outputSize);

  outputImage->SetLargestPossibleRegion(outputRegion);

  // Could be elsewhere...
  PointType outputOrigin;
  outputOrigin[0] = outputOrigin[1] = outputOrigin[2] = 0;

  outputImage->SetOrigin(outputOrigin);

  SpacingType inputSpacing = inputImage->GetSpacing();
  SpacingType outputSpacing;
  // Conserve axial spacing
  outputSpacing[2] = inputSpacing[m_ZDirection];
  // in-plane spacing = Radial spacing
  outputSpacing[0] = outputSpacing[1] = inputSpacing[m_RDirection];

  outputImage->SetSpacing(outputSpacing);
}

/**
 * Calculate necessary input image boundaries
 */
template <typename TInputImage, typename TOutputImage>
void
DirectFourierReconstructionImageToImageFilter<TInputImage, TOutputImage>::GenerateInputRequestedRegion()
{
  // call the superclass' implementation of this method
  Superclass::GenerateInputRequestedRegion();

  OutputImagePointer outputImage = this->GetOutput();
  InputImagePointer  inputImage = const_cast<InputImageType *>(this->GetInput());
  if (!inputImage || !outputImage)
  {
    return;
  }

  // request maximum angular and radial information / to be optimized

  SizeType  inputSize = inputImage->GetLargestPossibleRegion().GetSize();
  IndexType inputStart = inputImage->GetLargestPossibleRegion().GetIndex();

  // crop to requested z-axis
  inputSize[m_ZDirection] = outputImage->GetRequestedRegion().GetSize()[2];
  inputStart[m_ZDirection] = outputImage->GetRequestedRegion().GetIndex()[2];

  m_InputRequestedRegion.SetSize(inputSize);
  m_InputRequestedRegion.SetIndex(inputStart);

  m_InputRequestedRegion.Crop(inputImage->GetLargestPossibleRegion());
  inputImage->SetRequestedRegion(m_InputRequestedRegion);
}

/**
 * Actual computation
 */
template <typename TInputImage, typename TOutputImage>
void
DirectFourierReconstructionImageToImageFilter<TInputImage, TOutputImage>::GenerateData()
{
  OutputImagePointer     outputImage = this->GetOutput();
  ConstInputImagePointer inputImage = this->GetInput();

  if (!inputImage || !outputImage)
  {
    return;
  }

  outputImage->SetBufferedRegion(outputImage->GetRequestedRegion());
  outputImage->Allocate();

  // Crop angular input image size to 180 degrees
  typename InputImageType::RegionType inputROI = m_InputRequestedRegion;
  typename InputImageType::SizeType   inputROISize = inputROI.GetSize();
  typename InputImageType::IndexType  inputROIStart = inputROI.GetIndex();

  // the number of projections needed to cover 180 degrees
  const auto   alpha_size = Math::Floor<unsigned int>((180 * (inputROISize[m_AlphaDirection])) / m_AlphaRange);
  const double last_alpha_size = 1 + (180.0 * (inputROISize[m_AlphaDirection])) / m_AlphaRange - alpha_size;
  inputROIStart[m_AlphaDirection] += (inputROISize[m_AlphaDirection] - alpha_size) / 2;
  inputROISize[m_AlphaDirection] = alpha_size;
  inputROI.SetSize(inputROISize);
  inputROI.SetIndex(inputROIStart);

  // Setup the input ROI iterator
  InputSliceIteratorType inputIt(inputImage, inputROI);
  inputIt.SetFirstDirection(m_RDirection);      // Iterate first along r
  inputIt.SetSecondDirection(m_AlphaDirection); // then alpha (slice), and
                                                // finally z (stack)
  inputIt.GoToBegin();

  // Setup projection line
  ProjectionLineType::Pointer    projectionLine = ProjectionLineType::New();
  ProjectionLineType::RegionType pRegion;
  ProjectionLineType::SizeType   pSize;
  ProjectionLineType::IndexType  pStart;
  pSize[0] = inputROISize[m_RDirection] * m_ZeroPadding;
  pStart[0] = inputROIStart[m_RDirection];
  pRegion.SetSize(pSize);
  pRegion.SetIndex(pStart);
  projectionLine->SetRegions(pRegion);
  projectionLine->Allocate(true); // initialize
                                  // buffer to zero

  ProjectionLineType::IndexType pIdx;
  const unsigned int            pLineHalfShift = pSize[0] - inputROISize[m_RDirection] / 2;

  // Setup 1D FFT Filter
  FFTLineFilterType::Pointer FFT = FFTLineFilterType::New();
  FFT->SetInput(projectionLine);

  // Setup FFT Line interpolator stack
  auto * FFTLineInterpolator = new FFTLineInterpolatorType::Pointer[alpha_size];
  for (unsigned int alpha = 0; alpha < alpha_size; alpha++)
  {
    FFTLineInterpolator[alpha] = FFTLineInterpolatorType::New();
    FFTLineInterpolator[alpha]->SetSplineOrder(m_RadialSplineOrder);
  }

  // Setup cartesian FFTSlice domain
  FFTSliceType::RegionType FFTSliceRegion;
  FFTSliceType::SizeType   FFTSliceSize;
  FFTSliceType::IndexType  FFTSliceStart;

  FFTSliceSize[0] = FFTSliceSize[1] = inputROISize[m_RDirection] * m_ZeroPadding * m_OverSampling;
  FFTSliceStart[0] = FFTSliceStart[1] = 0;
  FFTSliceRegion.SetSize(FFTSliceSize);
  FFTSliceRegion.SetIndex(FFTSliceStart);

  FFTSliceType::Pointer FFTSlice = FFTSliceType::New();
  FFTSlice->SetRegions(FFTSliceRegion);
  FFTSlice->Allocate(true); // initialize
                            // buffer to zero

  FFTSliceIteratorType    FFTSliceIt(FFTSlice, FFTSliceRegion);
  FFTSliceType::IndexType sIdx;

  const double r_max = inputROISize[m_RDirection] * m_ZeroPadding * m_Cutoff / 2.0 - 1;
  const double halfR = inputROISize[m_RDirection] * m_ZeroPadding / 2;

  // Only the central frame will be copied to the output
  OutputSliceType::RegionType outputWindow;
  OutputSliceType::SizeType   outputWindowSize;
  OutputSliceType::IndexType  outputWindowStart;

  const unsigned int outputWindowShift = inputROISize[m_RDirection] * (m_ZeroPadding * m_OverSampling - 1) / 2;
  const unsigned int outputWindowHalf = inputROISize[m_RDirection] * m_ZeroPadding * m_OverSampling / 2;

  outputWindowSize[0] = outputImage->GetRequestedRegion().GetSize()[0];
  outputWindowSize[1] = outputImage->GetRequestedRegion().GetSize()[1];
  outputWindowStart[0] = outputImage->GetRequestedRegion().GetIndex()[0] + outputWindowShift;
  outputWindowStart[1] = outputImage->GetRequestedRegion().GetIndex()[1] + outputWindowShift;

  outputWindow.SetSize(outputWindowSize);
  outputWindow.SetIndex(outputWindowStart);

  OutputSliceType::IndexType          wIdx;
  typename OutputImageType::IndexType oIdx;

  ProgressReporter progress(this, 0, outputImage->GetRequestedRegion().GetNumberOfPixels());

  // Start iterating through slices
  while (!inputIt.IsAtEnd())
  {
    while (!inputIt.IsAtEndOfSlice()) // Start iterating through angles
    {
      while (!inputIt.IsAtEndOfLine()) // copy the whole input line
      {
        pIdx[0] = inputIt.GetIndex()[m_RDirection];

        // Shift the pixel to the borders of the image (origin @ 0)
        pIdx[0] += pLineHalfShift;
        pIdx[0] %= pSize[0];

        // Modulate image to shift DC to center of FFT line
        auto val = static_cast<ProjectionLineType::PixelType>((pIdx[0] & 1) ? -inputIt.Get() : inputIt.Get());
        projectionLine->SetPixel(pIdx, val);

        ++inputIt;
      } // while ( !inputIt.IsAtEndOfLine() )

      // Compute FFT
      FFT->Update();

      // link fft line into interpolator stack ...
      FFTLineInterpolator[inputIt.GetIndex()[m_AlphaDirection] - inputROIStart[m_AlphaDirection]]->SetInputImage(
        FFT->GetOutput());

      // ... and unlink from upstream pipeline
      FFT->GetOutput()->DisconnectPipeline();

      inputIt.NextLine();
    } // while ( !inputIt.IsAtEndOfSlice() )

    double       u, v;
    double       theta, r;
    double       alpha;
    unsigned int a_lo;

    // Resample the cartesian FFT Slice from polar lines
    for (FFTSliceIt.GoToBegin(); !FFTSliceIt.IsAtEnd(); ++FFTSliceIt)
    {
      sIdx = FFTSliceIt.GetIndex();

      // center on DC
      u = static_cast<double>(sIdx[0]) - static_cast<double>(FFTSliceSize[0]) / 2;
      v = static_cast<double>(sIdx[1]) - static_cast<double>(FFTSliceSize[1]) / 2;

      // Calculate polar radius
      r = sqrt(u * u + v * v) / m_OverSampling;

      // Radial cutoff frequency
      if (r >= r_max)
      {
        continue;
      }

      // Get polar angle - and map into [0 PI]
      if (u == 0.0 && v == 0.0)
      {
        theta = 0.0;
      }
      else
      {
        theta = std::atan2(v, u);
      }
      if (theta < 0)
      {
        theta += m_PI;
        r = -r;
      }

      // Convert into alpha-image indices
      alpha = theta * alpha_size / m_PI;
      if (alpha >= alpha_size)
      {
        alpha -= alpha_size;
        r = -r;
      }

      FFTLineType::PixelType out;

      // radial BSpline / linear angle interpolation
      a_lo = Math::Floor<unsigned int>(alpha);

      if (a_lo < alpha_size - 1) // no date-line crossing
      {
        // compute angular interpolation weights
        FFTLineType::PixelType w_a_lo(1 - (alpha - a_lo));
        FFTLineType::PixelType w_a_hi(alpha - a_lo);

        // get radial BSpline interpolations
        FFTLineInterpolatorType::ContinuousIndexType idx;
        idx[0] = r + halfR;
        FFTLineType::PixelType o_lo = FFTLineInterpolator[a_lo]->EvaluateAtContinuousIndex(idx);
        FFTLineType::PixelType o_hi = FFTLineInterpolator[a_lo + 1]->EvaluateAtContinuousIndex(idx);

        out = w_a_lo * o_lo + w_a_hi * o_hi;
      }
      else // date-line crossing
      {
        // compute angular interpolation weights (bigger last interval)
        FFTLineType::PixelType w_a_lo(1 - (alpha - a_lo) / last_alpha_size);
        FFTLineType::PixelType w_a_hi((alpha - a_lo) / last_alpha_size);

        // get radial BSpline interpolations
        FFTLineInterpolatorType::ContinuousIndexType idx;
        idx[0] = r + halfR;
        FFTLineType::PixelType o_lo = FFTLineInterpolator[a_lo]->EvaluateAtContinuousIndex(idx);
        idx[0] = halfR - r;
        FFTLineType::PixelType o_hi = FFTLineInterpolator[0]->EvaluateAtContinuousIndex(idx);

        out = w_a_lo * o_lo + w_a_hi * o_hi;
      }

      FFTSliceIt.Set(out);
    } // for FFTSliceIt

    // Setup inverse 2D FFT Filter
    IFFTSliceFilterType::Pointer IFFT = IFFTSliceFilterType::New();
    IFFT->SetInput(FFTSlice);

    // Calculate the inverse 2D FFT of the slice
    IFFT->Update();

    // Copy desired region into the outputImage

    // Setup output iterator
    OutputSliceIteratorType outputIt(IFFT->GetOutput(), outputWindow);
    // Current z-slice
    oIdx[2] = inputIt.GetIndex()[m_ZDirection];
    for (outputIt.GoToBegin(); !outputIt.IsAtEnd(); ++outputIt)
    {
      wIdx = outputIt.GetIndex();
      // Write to correct window
      oIdx[0] = wIdx[0] - outputWindowShift;
      oIdx[1] = wIdx[1] - outputWindowShift;
      // Fetch output data from buffer corners
      wIdx[0] += outputWindowHalf;
      wIdx[1] += outputWindowHalf;
      wIdx[0] %= FFTSliceSize[0];
      wIdx[1] %= FFTSliceSize[1];

      // Demodulate the image
      auto val = static_cast<OutputPixelType>(IFFT->GetOutput()->GetPixel(wIdx));
      outputImage->SetPixel(oIdx, ((wIdx[0] + wIdx[1]) & 1) ? -val : val);
      progress.CompletedPixel();
    } // for outputIt

    inputIt.NextSlice();
  } // while ( !inputIt.IsAtEnd() )

  delete[] FFTLineInterpolator;
}
} // namespace itk

#endif
