/*=========================================================================
 *
 *  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 itkBSplineControlPointImageFilter_hxx
#define itkBSplineControlPointImageFilter_hxx


#include "itkMath.h"
#include "itkImageDuplicator.h"
#include "itkImageRegionConstIteratorWithIndex.h"
#include "itkImageRegionIterator.h"
#include "itkImageRegionIteratorWithIndex.h"


namespace itk
{

template <typename TInputImage, typename TOutputImage>
BSplineControlPointImageFilter<TInputImage, TOutputImage>::BSplineControlPointImageFilter()
  : m_NumberOfLevels(1)

{
  this->m_Size.Fill(0);
  this->m_Spacing.Fill(1.0);
  this->m_Origin.Fill(0.0);
  this->m_Direction.SetIdentity();
  this->m_CloseDimension.Fill(0);
  this->m_SplineOrder.Fill(3);

  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    this->m_NumberOfControlPoints[i] = (this->m_SplineOrder[i] + 1);
    this->m_Kernel[i] = KernelType::New();
    this->m_Kernel[i]->SetSplineOrder(this->m_SplineOrder[i]);
  }
  this->m_KernelOrder0 = KernelOrder0Type::New();
  this->m_KernelOrder1 = KernelOrder1Type::New();
  this->m_KernelOrder2 = KernelOrder2Type::New();
  this->m_KernelOrder3 = KernelOrder3Type::New();
  this->DynamicMultiThreadingOn();
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::SetNumberOfLevels(ArrayType levels)
{
  this->m_NumberOfLevels = levels;
  this->m_MaximumNumberOfLevels = 1;
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    if (this->m_NumberOfLevels[i] == 0)
    {
      itkExceptionStringMacro("The number of levels in each dimension must be greater than 0");
    }
    if (this->m_NumberOfLevels[i] > this->m_MaximumNumberOfLevels)
    {
      this->m_MaximumNumberOfLevels = this->m_NumberOfLevels[i];
    }
  }

  itkDebugMacro("Setting m_NumberOfLevels to " << this->m_NumberOfLevels);
  itkDebugMacro("Setting m_MaximumNumberOfLevels to " << this->m_MaximumNumberOfLevels);

  this->m_DoMultilevel = this->m_MaximumNumberOfLevels > 1;
  this->SetSplineOrder(this->m_SplineOrder);
  this->Modified();
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::SetSplineOrder(unsigned int order)
{
  this->m_SplineOrder.Fill(order);
  this->SetSplineOrder(this->m_SplineOrder);
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::SetSplineOrder(ArrayType order)
{
  itkDebugMacro("Setting m_SplineOrder to " << order);

  this->m_SplineOrder = order;
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    if (this->m_SplineOrder[i] == 0)
    {
      itkExceptionStringMacro("The spline order in each dimension must be greater than 0");
    }

    this->m_Kernel[i] = KernelType::New();
    this->m_Kernel[i]->SetSplineOrder(this->m_SplineOrder[i]);

    if (this->m_DoMultilevel)
    {
      KernelType::MatrixType C = this->m_Kernel[i]->GetShapeFunctionsInZeroToOneInterval();

      vnl_matrix<RealType> R;
      vnl_matrix<RealType> S;
      R.set_size(C.rows(), C.cols());
      S.set_size(C.rows(), C.cols());
      for (unsigned int j = 0; j < C.rows(); ++j)
      {
        for (unsigned int k = 0; k < C.cols(); ++k)
        {
          R(j, k) = S(j, k) = static_cast<RealType>(C(j, k));
        }
      }
      for (unsigned int j = 0; j < C.cols(); ++j)
      {
        const RealType c = std::pow(static_cast<RealType>(2.0), static_cast<RealType>(C.cols() - j - 1));
        for (unsigned int k = 0; k < C.rows(); ++k)
        {
          R(k, j) *= c;
        }
      }
      R = R.transpose();
      R.flipud();
      S = S.transpose();
      S.flipud();

      this->m_RefinedLatticeCoefficients[i] = (vnl_svd<RealType>(R).solve(S)).extract(2, S.cols());
    }
  }
  this->Modified();
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::BeforeThreadedGenerateData()
{
  const TInputImage * inputPtr = this->GetInput();
  TOutputImage *      outputPtr = this->GetOutput();

  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    if (this->m_Size[i] == 0)
    {
      itkExceptionStringMacro("Size must be specified.");
    }
  }
  outputPtr->SetOrigin(this->m_Origin);
  outputPtr->SetSpacing(this->m_Spacing);
  outputPtr->SetRegions(this->m_Size);
  outputPtr->SetDirection(this->m_Direction);
  outputPtr->Allocate();

  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    this->m_NumberOfControlPoints[i] = inputPtr->GetLargestPossibleRegion().GetSize()[i];
  }
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::DynamicThreadedGenerateData(
  const OutputImageRegionType & region)
{
  const TInputImage * inputPtr = this->GetInput();
  TOutputImage *      outputPtr = this->GetOutput();

  typename PointDataImageType::Pointer collapsedPhiLattices[ImageDimension + 1];
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    collapsedPhiLattices[i] = PointDataImageType::New();
    collapsedPhiLattices[i]->CopyInformation(inputPtr);

    auto size = PointDataImageType::SizeType::Filled(1);
    for (unsigned int j = 0; j < i; ++j)
    {
      size[j] = inputPtr->GetLargestPossibleRegion().GetSize()[j];
    }
    collapsedPhiLattices[i]->SetRegions(size);
    collapsedPhiLattices[i]->Allocate();
  }
  using ImageDuplicatorType = ImageDuplicator<ControlPointLatticeType>;
  auto duplicator = ImageDuplicatorType::New();
  duplicator->SetInputImage(inputPtr);
  duplicator->Update();

  collapsedPhiLattices[ImageDimension] = duplicator->GetOutput();

  ArrayType totalNumberOfSpans;
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    if (this->m_CloseDimension[i])
    {
      totalNumberOfSpans[i] = inputPtr->GetLargestPossibleRegion().GetSize()[i];
    }
    else
    {
      totalNumberOfSpans[i] = inputPtr->GetLargestPossibleRegion().GetSize()[i] - this->m_SplineOrder[i];
    }
  }
  FixedArray<RealType, ImageDimension> U;
  auto                                 currentU = MakeFilled<FixedArray<RealType, ImageDimension>>(-1);

  typename OutputImageType::IndexType          startIndex = outputPtr->GetRequestedRegion().GetIndex();
  const typename PointDataImageType::IndexType startPhiIndex = inputPtr->GetLargestPossibleRegion().GetIndex();

  RealArrayType epsilon;
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    const RealType r = static_cast<RealType>(this->m_NumberOfControlPoints[i] - this->m_SplineOrder[i]) /
                       (static_cast<RealType>(this->m_Size[i] - 1) * this->m_Spacing[i]);
    epsilon[i] = r * this->m_Spacing[i] * this->m_BSplineEpsilon;
  }

  for (ImageRegionIteratorWithIndex<OutputImageType> It(outputPtr, region); !It.IsAtEnd(); ++It)
  {
    typename OutputImageType::IndexType idx = It.GetIndex();
    for (unsigned int i = 0; i < ImageDimension; ++i)
    {
      U[i] = static_cast<RealType>(totalNumberOfSpans[i]) * static_cast<RealType>(idx[i] - startIndex[i]) /
             static_cast<RealType>(this->m_Size[i] - 1);
      if (itk::Math::Absolute(U[i] - static_cast<RealType>(totalNumberOfSpans[i])) <= epsilon[i])
      {
        U[i] = static_cast<RealType>(totalNumberOfSpans[i]) - epsilon[i];
      }
      if (U[i] < RealType{} && itk::Math::Absolute(U[i]) <= epsilon[i])
      {
        U[i] = RealType{};
      }

      if (U[i] < RealType{} || U[i] >= static_cast<RealType>(totalNumberOfSpans[i]))
      {
        itkExceptionMacro("The collapse point component "
                          << U[i] << " is outside the corresponding parametric domain of [0, " << totalNumberOfSpans[i]
                          << ").");
      }
    }
    for (int i = ImageDimension - 1; i >= 0; i--)
    {
      if (Math::NotExactlyEquals(U[i], currentU[i]))
      {
        for (int j = i; j >= 0; j--)
        {
          this->CollapsePhiLattice(collapsedPhiLattices[j + 1], collapsedPhiLattices[j], U[j], j);
          currentU[j] = U[j];
        }
        break;
      }
    }
    It.Set(collapsedPhiLattices[0]->GetPixel(startPhiIndex));
  }
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::CollapsePhiLattice(PointDataImageType * lattice,
                                                                              PointDataImageType * collapsedLattice,
                                                                              const RealType       u,
                                                                              const unsigned int   dimension)
{
  for (ImageRegionIteratorWithIndex<PointDataImageType> It(collapsedLattice,
                                                           collapsedLattice->GetLargestPossibleRegion());
       !It.IsAtEnd();
       ++It)
  {
    PointDataType                          data{};
    typename PointDataImageType::IndexType idx = It.GetIndex();
    for (unsigned int i = 0; i < this->m_SplineOrder[dimension] + 1; ++i)
    {
      idx[dimension] = static_cast<unsigned int>(u) + i;
      const RealType v = u - idx[dimension] + 0.5 * static_cast<RealType>(this->m_SplineOrder[dimension] - 1);

      RealType B = 0.0;
      switch (this->m_SplineOrder[dimension])
      {
        case 0:
        {
          B = this->m_KernelOrder0->Evaluate(v);
          break;
        }
        case 1:
        {
          B = this->m_KernelOrder1->Evaluate(v);
          break;
        }
        case 2:
        {
          B = this->m_KernelOrder2->Evaluate(v);
          break;
        }
        case 3:
        {
          B = this->m_KernelOrder3->Evaluate(v);
          break;
        }
        default:
        {
          B = this->m_Kernel[dimension]->Evaluate(v);
          break;
        }
      }
      if (this->m_CloseDimension[dimension])
      {
        idx[dimension] %= lattice->GetLargestPossibleRegion().GetSize()[dimension];
      }
      data += (lattice->GetPixel(idx) * B);
    }
    It.Set(data);
  }
}

template <typename TInputImage, typename TOutputImage>
unsigned int
BSplineControlPointImageFilter<TInputImage, TOutputImage>::SplitRequestedRegion(unsigned int            i,
                                                                                unsigned int            num,
                                                                                OutputImageRegionType & splitRegion)
{
  // Get the output pointer
  OutputImageType * outputPtr = this->GetOutput();

  const SizeType requestedRegionSize = outputPtr->GetRequestedRegion().GetSize();

  // Initialize the splitRegion to the output requested region
  splitRegion = outputPtr->GetRequestedRegion();
  typename TOutputImage::IndexType splitIndex = splitRegion.GetIndex();
  typename TOutputImage::SizeType  splitSize = splitRegion.GetSize();

  // split on the outermost dimension
  const auto splitAxis = outputPtr->GetImageDimension() - 1;

  // determine the actual number of pieces that will be generated
  const typename SizeType::SizeValueType range = requestedRegionSize[splitAxis];
  auto valuesPerThread = static_cast<unsigned int>(std::ceil(range / static_cast<double>(num)));
  auto maxThreadIdUsed = static_cast<unsigned int>(std::ceil(range / static_cast<double>(valuesPerThread)) - 1);

  // Split the region
  if (i < maxThreadIdUsed)
  {
    splitIndex[splitAxis] += i * valuesPerThread;
    splitSize[splitAxis] = valuesPerThread;
  }
  if (i == maxThreadIdUsed)
  {
    splitIndex[splitAxis] += i * valuesPerThread;
    // last thread needs to process the "rest" dimension being split
    splitSize[splitAxis] = splitSize[splitAxis] - i * valuesPerThread;
  }

  // set the split region ivars
  splitRegion = { splitIndex, splitSize };

  itkDebugMacro("Split piece: " << splitRegion);

  return maxThreadIdUsed + 1;
}

template <typename TInputPointImage, typename TOutputImage>
auto
BSplineControlPointImageFilter<TInputPointImage, TOutputImage>::RefineControlPointLattice(ArrayType numberOfLevels) ->
  typename ControlPointLatticeType::Pointer
{
  this->SetNumberOfLevels(numberOfLevels);

  using ImageDuplicatorType = ImageDuplicator<ControlPointLatticeType>;
  auto duplicator = ImageDuplicatorType::New();
  duplicator->SetInputImage(this->GetInput());
  duplicator->Update();

  auto psiLattice = ControlPointLatticeType::New();
  psiLattice = duplicator->GetOutput();

  for (unsigned int m = 1; m < this->m_MaximumNumberOfLevels; ++m)
  {
    ArrayType numberOfNewControlPoints;
    for (unsigned int i = 0; i < ImageDimension; ++i)
    {
      numberOfNewControlPoints[i] = psiLattice->GetLargestPossibleRegion().GetSize()[i];
    }
    for (unsigned int i = 0; i < ImageDimension; ++i)
    {
      if (m < this->m_NumberOfLevels[i])
      {
        numberOfNewControlPoints[i] = 2 * numberOfNewControlPoints[i] - this->m_SplineOrder[i];
      }
    }
    typename RealImageType::RegionType::SizeType size;
    for (unsigned int i = 0; i < ImageDimension; ++i)
    {
      if (this->m_CloseDimension[i])
      {
        size[i] = numberOfNewControlPoints[i] - this->m_SplineOrder[i];
      }
      else
      {
        size[i] = numberOfNewControlPoints[i];
      }
    }

    auto refinedLattice = ControlPointLatticeType::New();
    refinedLattice->SetRegions(size);
    refinedLattice->Allocate();
    constexpr PixelType data{};
    refinedLattice->FillBuffer(data);

    typename ControlPointLatticeType::IndexType            idx;
    typename ControlPointLatticeType::IndexType            idxPsi;
    typename ControlPointLatticeType::IndexType            tmp;
    typename ControlPointLatticeType::IndexType            tmpPsi;
    typename ControlPointLatticeType::IndexType            off;
    typename ControlPointLatticeType::IndexType            offPsi;
    typename ControlPointLatticeType::RegionType::SizeType sizePsi;

    size.Fill(2);
    unsigned int N = 1;
    for (unsigned int i = 0; i < ImageDimension; ++i)
    {
      N *= (this->m_SplineOrder[i] + 1);
      sizePsi[i] = this->m_SplineOrder[i] + 1;
    }

    ImageRegionIteratorWithIndex<ControlPointLatticeType> It(refinedLattice,
                                                             refinedLattice->GetLargestPossibleRegion());

    const TInputPointImage * input = this->GetInput();

    It.GoToBegin();
    while (!It.IsAtEnd())
    {
      idx = It.GetIndex();
      for (unsigned int i = 0; i < ImageDimension; ++i)
      {
        if (m < this->m_NumberOfLevels[i])
        {
          idxPsi[i] = static_cast<unsigned int>(0.5 * idx[i]);
        }
        else
        {
          idxPsi[i] = static_cast<unsigned int>(idx[i]);
        }
      }
      for (unsigned int i = 0; i < (2 << (ImageDimension - 1)); ++i)
      {
        PixelType sum{};

        PixelType val;
        off = this->NumberToIndex(i, size);

        bool outOfBoundary = false;
        for (unsigned int j = 0; j < ImageDimension; ++j)
        {
          tmp[j] = idx[j] + off[j];
          if (tmp[j] >= static_cast<int>(numberOfNewControlPoints[j]) && !this->m_CloseDimension[j])
          {
            outOfBoundary = true;
            break;
          }
          if (this->m_CloseDimension[j])
          {
            tmp[j] %= refinedLattice->GetLargestPossibleRegion().GetSize()[j];
          }
        }
        if (outOfBoundary)
        {
          continue;
        }

        for (unsigned int j = 0; j < N; ++j)
        {
          offPsi = this->NumberToIndex(j, sizePsi);

          bool outOfBoundary2 = false;
          for (unsigned int k = 0; k < ImageDimension; ++k)
          {
            tmpPsi[k] = idxPsi[k] + offPsi[k];
            if (tmpPsi[k] >= static_cast<int>(input->GetLargestPossibleRegion().GetSize()[k]) &&
                !this->m_CloseDimension[k])
            {
              outOfBoundary2 = true;
              break;
            }
            if (this->m_CloseDimension[k])
            {
              tmpPsi[k] %= psiLattice->GetLargestPossibleRegion().GetSize()[k];
            }
          }
          if (outOfBoundary2)
          {
            continue;
          }
          RealType coeff = 1.0;
          for (unsigned int k = 0; k < ImageDimension; ++k)
          {
            coeff *= this->m_RefinedLatticeCoefficients[k](off[k], offPsi[k]);
          }
          val = psiLattice->GetPixel(tmpPsi);
          val *= coeff;
          sum += val;
        }
        refinedLattice->SetPixel(tmp, sum);
      }

      bool IsEvenIndex = false;
      while (!IsEvenIndex && !It.IsAtEnd())
      {
        ++It;
        idx = It.GetIndex();
        IsEvenIndex = true;
        for (unsigned int i = 0; i < ImageDimension; ++i)
        {
          if (idx[i] % 2)
          {
            IsEvenIndex = false;
          }
        }
      }
    }
    auto duplicator2 = ImageDuplicatorType::New();
    duplicator2->SetInputImage(refinedLattice);
    duplicator2->Update();
    psiLattice = duplicator2->GetOutput();
  }

  // Specify the pose parameters of the control point lattice

  typename PointDataImageType::PointType   origin;
  typename PointDataImageType::SpacingType spacing;

  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    const RealType domain = this->m_Spacing[i] * static_cast<RealType>(this->m_Size[i] - 1);

    unsigned int totalNumberOfSpans = psiLattice->GetLargestPossibleRegion().GetSize()[i];
    if (!this->m_CloseDimension[i])
    {
      totalNumberOfSpans -= this->m_SplineOrder[i];
    }

    spacing[i] = domain / static_cast<RealType>(totalNumberOfSpans);

    origin[i] = -0.5 * spacing[i] * (this->m_SplineOrder[i] - 1);
  }
  origin = this->m_Direction * origin;

  psiLattice->SetOrigin(origin);
  psiLattice->SetSpacing(spacing);
  psiLattice->SetDirection(this->m_Direction);

  return psiLattice;
}

template <typename TInputImage, typename TOutputImage>
void
BSplineControlPointImageFilter<TInputImage, TOutputImage>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    this->m_Kernel[i]->Print(os, indent.GetNextIndent());
  }
  os << indent << "Spline order: " << this->m_SplineOrder << std::endl;
  os << indent << "Close dimension: " << this->m_CloseDimension << std::endl;
  os << indent << "Parametric domain" << std::endl;
  os << indent << "  Origin:    " << this->m_Origin << std::endl;
  os << indent << "  Spacing:   " << this->m_Spacing << std::endl;
  os << indent << "  Size:      " << this->m_Size << std::endl;
  os << indent << "  Direction: " << this->m_Direction << std::endl;
}

} // end namespace itk

#endif
