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

#include "itkProgressReporter.h"
#include "itkShapedNeighborhoodIterator.h"
#include "itkImageRegionIterator.h"
#include "itkConstantBoundaryCondition.h"
#include "itkImageRegionIteratorWithIndex.h"
#include "itkImageRegionIterator.h"
#include "itkOffset.h"
#include "itkNeighborhoodAlgorithm.h"
#include "itkZeroFluxNeumannBoundaryCondition.h"
#include <algorithm> // For min.

/*
 *
 * This code was contributed in the Insight Journal paper:
 * "Efficient implementation of kernel filtering"
 * by Beare R., Lehmann G
 * https://www.insight-journal.org/browse/publication/160
 *
 */

namespace itk_impl_details
{

template <typename TIterator>
inline TIterator *
setConnectivityEarlyBox(TIterator * it, bool fullyConnected = false)
{
  // activate the "previous" neighbours
  typename TIterator::OffsetType offset;
  it->ClearActiveList();
  if (!fullyConnected)
  {
    // only activate the neighbors that are face connected
    // to the current pixel. do not include the center pixel
    offset.Fill(0);
    for (unsigned int d = 0; d < TIterator::Dimension; ++d)
    {
      offset[d] = -1;
      it->ActivateOffset(offset);
      offset[d] = 0;
    }
  }
  else
  {
    // activate all neighbors that are face+edge+vertex
    // connected to the current pixel. do not include the center pixel
    unsigned int centerIndex = it->GetCenterNeighborhoodIndex();
    for (unsigned int d = 0; d < centerIndex; ++d)
    {
      offset = it->GetOffset(d);
      // check for positives in any dimension
      bool keep = true;
      for (unsigned int i = 0; i < TIterator::Dimension; ++i)
      {
        if (offset[i] > 0)
        {
          keep = false;
          break;
        }
      }
      if (keep)
      {
        it->ActivateOffset(offset);
      }
    }
    offset.Fill(0);
    it->DeactivateOffset(offset);
  }
  return it;
}

} // namespace itk_impl_details

namespace itk
{

template <typename TInputImage, typename TOutputImage>
void
BoxAccumulateFunction(const TInputImage *               inputImage,
                      const TOutputImage *              outputImage,
                      typename TInputImage::RegionType  inputRegion,
                      typename TOutputImage::RegionType outputRegion
#if defined(ITKV4_COMPATIBILITY)
                      ,
                      ProgressReporter & progress)
#else
)
#endif
{
  // type alias
  using InputImageType = TInputImage;
  using OffsetType = typename TInputImage::OffsetType;
  using OutputImageType = TOutputImage;
  using OutputPixelType = typename TOutputImage::PixelType;

  using InputIterator = ImageRegionConstIterator<TInputImage>;

  using NOutputIterator = ShapedNeighborhoodIterator<TOutputImage>;
  InputIterator                  inIt(inputImage, inputRegion);
  typename TInputImage::SizeType kernelRadius;
  kernelRadius.Fill(1);

  NOutputIterator noutIt(kernelRadius, outputImage, outputRegion);
  // this iterator is fully connected
  itk_impl_details::setConnectivityEarlyBox(&noutIt, true);

  ConstantBoundaryCondition<OutputImageType> oBC;
  oBC.SetConstant(OutputPixelType{});
  noutIt.OverrideBoundaryCondition(&oBC);
  // This uses several iterators. An alternative and probably better
  // approach would be to copy the input to the output and convolve
  // with the following weights (in 2D)
  //   -(dim - 1)  1
  //       1       1
  // The result of each convolution needs to get written back to the
  // image being convolved so that the accumulation propagates
  // This should be implementable with neighborhood operators.

  std::vector<int>                        weights;
  typename NOutputIterator::ConstIterator sIt;
  for (auto idxIt = noutIt.GetActiveIndexList().begin(); idxIt != noutIt.GetActiveIndexList().end(); ++idxIt)
  {
    OffsetType offset = noutIt.GetOffset(*idxIt);
    int        w = -1;
    for (unsigned int k = 0; k < InputImageType::ImageDimension; ++k)
    {
      if (offset[k] != 0)
      {
        w *= offset[k];
      }
    }
    //     std::cout << offset << "  " << w << std::endl;
    weights.push_back(w);
  }

  for (inIt.GoToBegin(), noutIt.GoToBegin(); !noutIt.IsAtEnd(); ++inIt, ++noutIt)
  {
    OutputPixelType sum = 0;
    int             k;
    for (k = 0, sIt = noutIt.Begin(); !sIt.IsAtEnd(); ++sIt, ++k)
    {
      sum += sIt.Get() * weights[k];
    }
    noutIt.SetCenterPixel(sum + inIt.Get());
#if defined(ITKV4_COMPATIBILITY)
    progress.CompletedPixel();
#endif
  }
}

// a function to generate corners of arbitrary dimension box
template <typename TImage>
std::vector<typename TImage::OffsetType>
CornerOffsets(const TImage * im)
{
  using NIterator = ShapedNeighborhoodIterator<TImage>;
  typename TImage::SizeType unitradius;
  unitradius.Fill(1);
  NIterator                                n1(unitradius, im, im->GetRequestedRegion());
  unsigned int                             centerIndex = n1.GetCenterNeighborhoodIndex();
  typename NIterator::OffsetType           offset;
  std::vector<typename TImage::OffsetType> result;
  for (unsigned int d = 0; d < centerIndex * 2 + 1; ++d)
  {
    offset = n1.GetOffset(d);
    // check whether this is a corner - corners have no zeros
    bool corner = true;
    for (unsigned int k = 0; k < TImage::ImageDimension; ++k)
    {
      if (offset[k] == 0)
      {
        corner = false;
        break;
      }
    }
    if (corner)
    {
      result.push_back(offset);
    }
  }
  return (result);
}

template <typename TInputImage, typename TOutputImage>
void
BoxMeanCalculatorFunction(const TInputImage *               accImage,
                          TOutputImage *                    outputImage,
                          typename TInputImage::RegionType  inputRegion,
                          typename TOutputImage::RegionType outputRegion,
                          typename TInputImage::SizeType    radius
#if defined(ITKV4_COMPATIBILITY)
                          ,
                          ProgressReporter & progress)
#else
)
#endif
{
  // type alias
  using InputImageType = TInputImage;
  using RegionType = typename TInputImage::RegionType;
  using SizeType = typename TInputImage::SizeType;
  using IndexType = typename TInputImage::IndexType;
  using OffsetType = typename TInputImage::OffsetType;
  using OutputImageType = TOutputImage;
  using OutputPixelType = typename TOutputImage::PixelType;
  // use the face generator for speed
  using FaceCalculatorType = NeighborhoodAlgorithm::ImageBoundaryFacesCalculator<InputImageType>;
  using FaceListType = typename FaceCalculatorType::FaceListType;
  FaceCalculatorType faceCalculator;

  ZeroFluxNeumannBoundaryCondition<TInputImage> nbc;

  // this process is actually slightly asymmetric because we need to
  // subtract rectangles that are next to our kernel, not overlapping it
  SizeType kernelSize;
  SizeType internalRadius;
  SizeType regionLimit;

  IndexType regionStart = inputRegion.GetIndex();
  for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
  {
    kernelSize[i] = radius[i] * 2 + 1;
    internalRadius[i] = radius[i] + 1;
    regionLimit[i] = inputRegion.GetSize()[i] + regionStart[i] - 1;
  }

  using AccPixType = typename NumericTraits<OutputPixelType>::RealType;
  // get a set of offsets to corners for a unit hypercube in this image
  std::vector<OffsetType> unitCorners = CornerOffsets<TInputImage>(accImage);
  std::vector<OffsetType> realCorners;
  std::vector<AccPixType> weights;
  // now compute the weights
  for (unsigned int k = 0; k < unitCorners.size(); ++k)
  {
    int        prod = 1;
    OffsetType thisCorner;
    for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
    {
      prod *= unitCorners[k][i];
      if (unitCorners[k][i] > 0)
      {
        thisCorner[i] = radius[i];
      }
      else
      {
        thisCorner[i] = -(static_cast<OffsetValueType>(radius[i]) + 1);
      }
    }
    weights.push_back((AccPixType)prod);
    realCorners.push_back(thisCorner);
  }

  FaceListType faceList = faceCalculator(accImage, outputRegion, internalRadius);
  // start with the body region
  for (const auto & face : faceList)
  {
    if (&face == &faceList.front())
    {
      // this is the body region. This is meant to be an optimized
      // version that doesn't use neighborhood regions
      // compute the various offsets
      AccPixType pixelscount = 1;
      for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
      {
        pixelscount *= (AccPixType)(2 * radius[i] + 1);
      }

      using OutputIteratorType = ImageRegionIterator<OutputImageType>;
      using InputIteratorType = ImageRegionConstIterator<InputImageType>;

      using CornerItVecType = std::vector<InputIteratorType>;
      CornerItVecType cornerItVec;
      // set up the iterators for each corner
      for (unsigned int k = 0; k < realCorners.size(); ++k)
      {
        typename InputImageType::RegionType tReg = face;
        tReg.SetIndex(tReg.GetIndex() + realCorners[k]);
        InputIteratorType tempIt(accImage, tReg);
        tempIt.GoToBegin();
        cornerItVec.push_back(tempIt);
      }
      // set up the output iterator
      OutputIteratorType oIt(outputImage, face);
      // now do the work
      for (oIt.GoToBegin(); !oIt.IsAtEnd(); ++oIt)
      {
        AccPixType sum = 0;
        // check each corner
        for (unsigned int k = 0; k < cornerItVec.size(); ++k)
        {
          sum += weights[k] * cornerItVec[k].Get();
          // increment each corner iterator
          ++(cornerItVec[k]);
        }
        oIt.Set(static_cast<OutputPixelType>(sum / pixelscount));
#if defined(ITKV4_COMPATIBILITY)
        progress.CompletedPixel();
#endif
      }
    }
    else
    {
      // now we need to deal with the border regions
      using OutputIteratorType = ImageRegionIteratorWithIndex<OutputImageType>;
      OutputIteratorType oIt(outputImage, face);
      // now do the work
      for (oIt.GoToBegin(); !oIt.IsAtEnd(); ++oIt)
      {
        // figure out the number of pixels in the box by creating an
        // equivalent region and cropping - this could probably be
        // included in the loop below.
        RegionType currentKernelRegion;
        currentKernelRegion.SetSize(kernelSize);
        // compute the region's index
        IndexType kernelRegionIdx = oIt.GetIndex();
        IndexType centIndex = kernelRegionIdx;
        for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
        {
          kernelRegionIdx[i] -= radius[i];
        }
        currentKernelRegion.SetIndex(kernelRegionIdx);
        currentKernelRegion.Crop(inputRegion);
        OffsetValueType edgepixelscount = currentKernelRegion.GetNumberOfPixels();
        AccPixType      sum = 0;
        // rules are : for each corner,
        //               for each dimension
        //                  if dimension offset is positive -> this is
        //                  a leading edge. Crop if outside the input
        //                  region
        //                  if dimension offset is negative -> this is
        //                  a trailing edge. Ignore if it is outside
        //                  image region
        for (unsigned int k = 0; k < realCorners.size(); ++k)
        {
          IndexType thisCorner = centIndex + realCorners[k];
          bool      includeCorner = true;
          for (unsigned int j = 0; j < TInputImage::ImageDimension; ++j)
          {
            if (unitCorners[k][j] > 0)
            {
              // leading edge - crop it
              thisCorner[j] = std::min(thisCorner[j], static_cast<OffsetValueType>(regionLimit[j]));
            }
            else
            {
              // trailing edge - check bounds
              if (thisCorner[j] < regionStart[j])
              {
                includeCorner = false;
                break;
              }
            }
          }
          if (includeCorner)
          {
            sum += accImage->GetPixel(thisCorner) * weights[k];
          }
        }

        oIt.Set(static_cast<OutputPixelType>(sum / (AccPixType)edgepixelscount));
#if defined(ITKV4_COMPATIBILITY)
        progress.CompletedPixel();
#endif
      }
    }
  }
}

template <typename TInputImage, typename TOutputImage>
void
BoxSigmaCalculatorFunction(const TInputImage *               accImage,
                           TOutputImage *                    outputImage,
                           typename TInputImage::RegionType  inputRegion,
                           typename TOutputImage::RegionType outputRegion,
                           typename TInputImage::SizeType    radius
#if defined(ITKV4_COMPATIBILITY)
                           ,
                           ProgressReporter & progress)
#else
)
#endif
{
  // type alias
  using InputImageType = TInputImage;
  using RegionType = typename TInputImage::RegionType;
  using SizeType = typename TInputImage::SizeType;
  using IndexType = typename TInputImage::IndexType;
  using OffsetType = typename TInputImage::OffsetType;
  using OutputImageType = TOutputImage;
  using OutputPixelType = typename TOutputImage::PixelType;
  using InputPixelType = typename TInputImage::PixelType;
  // use the face generator for speed
  using FaceCalculatorType = typename NeighborhoodAlgorithm::ImageBoundaryFacesCalculator<InputImageType>;
  using FaceListType = typename FaceCalculatorType::FaceListType;
  FaceCalculatorType faceCalculator;

  ZeroFluxNeumannBoundaryCondition<TInputImage> nbc;

  // this process is actually slightly asymmetric because we need to
  // subtract rectangles that are next to our kernel, not overlapping it
  SizeType  kernelSize;
  SizeType  internalRadius;
  SizeType  regionLimit;
  IndexType regionStart = inputRegion.GetIndex();
  for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
  {
    kernelSize[i] = radius[i] * 2 + 1;
    internalRadius[i] = radius[i] + 1;
    regionLimit[i] = inputRegion.GetSize()[i] + regionStart[i] - 1;
  }

  using AccPixType = typename NumericTraits<OutputPixelType>::RealType;
  // get a set of offsets to corners for a unit hypercube in this image
  std::vector<OffsetType> unitCorners = CornerOffsets<TInputImage>(accImage);
  std::vector<OffsetType> realCorners;
  std::vector<AccPixType> weights;
  // now compute the weights
  for (unsigned int k = 0; k < unitCorners.size(); ++k)
  {
    int        prod = 1;
    OffsetType thisCorner;
    for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
    {
      prod *= unitCorners[k][i];
      if (unitCorners[k][i] > 0)
      {
        thisCorner[i] = radius[i];
      }
      else
      {
        thisCorner[i] = -(static_cast<OffsetValueType>(radius[i]) + 1);
      }
    }
    weights.push_back((AccPixType)prod);
    realCorners.push_back(thisCorner);
  }

  FaceListType faceList = faceCalculator(accImage, outputRegion, internalRadius);
  // start with the body region
  for (const auto & face : faceList)
  {
    if (&face == &faceList.front())
    {
      // this is the body region. This is meant to be an optimized
      // version that doesn't use neighborhood regions
      // compute the various offsets
      AccPixType pixelscount = 1;
      for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
      {
        pixelscount *= (AccPixType)(2 * radius[i] + 1);
      }

      using OutputIteratorType = ImageRegionIterator<OutputImageType>;
      using InputIteratorType = ImageRegionConstIterator<InputImageType>;

      using CornerItVecType = std::vector<InputIteratorType>;
      CornerItVecType cornerItVec;
      // set up the iterators for each corner
      for (unsigned int k = 0; k < realCorners.size(); ++k)
      {
        typename InputImageType::RegionType tReg = face;
        tReg.SetIndex(tReg.GetIndex() + realCorners[k]);
        InputIteratorType tempIt(accImage, tReg);
        tempIt.GoToBegin();
        cornerItVec.push_back(tempIt);
      }
      // set up the output iterator
      OutputIteratorType oIt(outputImage, face);
      // now do the work
      for (oIt.GoToBegin(); !oIt.IsAtEnd(); ++oIt)
      {
        AccPixType sum = 0;
        AccPixType squareSum = 0;
        // check each corner
        for (unsigned int k = 0; k < cornerItVec.size(); ++k)
        {
          const InputPixelType & i = cornerItVec[k].Get();
          sum += weights[k] * i[0];
          squareSum += weights[k] * i[1];
          // increment each corner iterator
          ++(cornerItVec[k]);
        }

        oIt.Set(static_cast<OutputPixelType>(std::sqrt((squareSum - sum * sum / pixelscount) / (pixelscount - 1))));
#if defined(ITKV4_COMPATIBILITY)
        progress.CompletedPixel();
#endif
      }
    }
    else
    {
      // now we need to deal with the border regions
      using OutputIteratorType = ImageRegionIteratorWithIndex<OutputImageType>;
      OutputIteratorType oIt(outputImage, face);
      // now do the work
      for (oIt.GoToBegin(); !oIt.IsAtEnd(); ++oIt)
      {
        // figure out the number of pixels in the box by creating an
        // equivalent region and cropping - this could probably be
        // included in the loop below.
        RegionType currentKernelRegion;
        currentKernelRegion.SetSize(kernelSize);
        // compute the region's index
        IndexType kernelRegionIdx = oIt.GetIndex();
        IndexType centIndex = kernelRegionIdx;
        for (unsigned int i = 0; i < TInputImage::ImageDimension; ++i)
        {
          kernelRegionIdx[i] -= radius[i];
        }
        currentKernelRegion.SetIndex(kernelRegionIdx);
        currentKernelRegion.Crop(inputRegion);
        SizeValueType edgepixelscount = currentKernelRegion.GetNumberOfPixels();
        AccPixType    sum = 0;
        AccPixType    squareSum = 0;
        // rules are : for each corner,
        //               for each dimension
        //                  if dimension offset is positive -> this is
        //                  a leading edge. Crop if outside the input
        //                  region
        //                  if dimension offset is negative -> this is
        //                  a trailing edge. Ignore if it is outside
        //                  image region
        for (unsigned int k = 0; k < realCorners.size(); ++k)
        {
          IndexType thisCorner = centIndex + realCorners[k];
          bool      includeCorner = true;
          for (unsigned int j = 0; j < TInputImage::ImageDimension; ++j)
          {
            if (unitCorners[k][j] > 0)
            {
              // leading edge - crop it
              thisCorner[j] = std::min(thisCorner[j], static_cast<OffsetValueType>(regionLimit[j]));
            }
            else
            {
              // trailing edge - check bounds
              if (thisCorner[j] < regionStart[j])
              {
                includeCorner = false;
                break;
              }
            }
          }
          if (includeCorner)
          {
            const InputPixelType & i = accImage->GetPixel(thisCorner);
            sum += weights[k] * i[0];
            squareSum += weights[k] * i[1];
          }
        }

        oIt.Set(
          static_cast<OutputPixelType>(std::sqrt((squareSum - sum * sum / edgepixelscount) / (edgepixelscount - 1))));
#if defined(ITKV4_COMPATIBILITY)
        progress.CompletedPixel();
#endif
      }
    }
  }
}

template <typename TInputImage, typename TOutputImage>
void
BoxSquareAccumulateFunction(const TInputImage *               inputImage,
                            TOutputImage *                    outputImage,
                            typename TInputImage::RegionType  inputRegion,
                            typename TOutputImage::RegionType outputRegion
#if defined(ITKV4_COMPATIBILITY)
                            ,
                            ProgressReporter & progress)
#else
)
#endif
{
  // type alias
  using InputImageType = TInputImage;
  using OffsetType = typename TInputImage::OffsetType;
  using OutputImageType = TOutputImage;
  using OutputPixelType = typename TOutputImage::PixelType;
  using ValueType = typename OutputPixelType::ValueType;
  using InputPixelType = typename TInputImage::PixelType;

  using InputIterator = ImageRegionConstIterator<TInputImage>;

  using NOutputIterator = ShapedNeighborhoodIterator<TOutputImage>;
  InputIterator                  inIt(inputImage, inputRegion);
  typename TInputImage::SizeType kernelRadius;
  kernelRadius.Fill(1);

  NOutputIterator noutIt(kernelRadius, outputImage, outputRegion);
  // this iterator is fully connected
  itk_impl_details::setConnectivityEarlyBox(&noutIt, true);

  ConstantBoundaryCondition<OutputImageType> oBC;
  oBC.SetConstant(OutputPixelType{});
  noutIt.OverrideBoundaryCondition(&oBC);
  // This uses several iterators. An alternative and probably better
  // approach would be to copy the input to the output and convolve
  // with the following weights (in 2D)
  //   -(dim - 1)  1
  //       1       1
  // The result of each convolution needs to get written back to the
  // image being convolved so that the accumulation propagates
  // This should be implementable with neighborhood operators.

  std::vector<int>                        weights;
  typename NOutputIterator::ConstIterator sIt;
  for (auto idxIt = noutIt.GetActiveIndexList().begin(); idxIt != noutIt.GetActiveIndexList().end(); ++idxIt)
  {
    OffsetType offset = noutIt.GetOffset(*idxIt);
    int        w = -1;
    for (unsigned int k = 0; k < InputImageType::ImageDimension; ++k)
    {
      if (offset[k] != 0)
      {
        w *= offset[k];
      }
    }
    weights.push_back(w);
  }

  for (inIt.GoToBegin(), noutIt.GoToBegin(); !noutIt.IsAtEnd(); ++inIt, ++noutIt)
  {
    ValueType sum = 0;
    ValueType squareSum = 0;
    int       k;
    for (k = 0, sIt = noutIt.Begin(); !sIt.IsAtEnd(); ++sIt, ++k)
    {
      const OutputPixelType & v = sIt.Get();
      sum += v[0] * weights[k];
      squareSum += v[1] * weights[k];
    }
    OutputPixelType        o;
    const InputPixelType & i = inIt.Get();
    o[0] = sum + i;
    o[1] = squareSum + i * i;
    noutIt.SetCenterPixel(o);
#if defined(ITKV4_COMPATIBILITY)
    progress.CompletedPixel();
#endif
  }
}
} // namespace itk

#endif