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

#include "itkImageRegionIteratorWithIndex.h"
#include "itkOffset.h"
#include "itkTotalProgressReporter.h"
#include "itkNumericTraits.h"

#include "itkImageRegionIterator.h"
#include "itkImageLinearConstIteratorWithIndex.h"

/*
 *
 * 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
{
template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::
  MaskedMovingHistogramImageFilter()
{
  this->SetNumberOfRequiredInputs(2);
  this->m_FillValue = OutputPixelType{};
  this->m_MaskValue = NumericTraits<MaskPixelType>::max();
  this->m_BackgroundMaskValue = MaskPixelType{};
  this->m_GenerateOutputMask = true;
  this->SetGenerateOutputMask(false);
  this->DynamicMultiThreadingOn();
  this->ThreaderUpdateProgressOff();
}

template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
void
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::SetGenerateOutputMask(
  bool generateOutputMask)
{
  if (generateOutputMask != this->m_GenerateOutputMask)
  {
    this->m_GenerateOutputMask = generateOutputMask;
    if (generateOutputMask)
    {
      this->SetNumberOfRequiredOutputs(2);
      typename MaskImageType::Pointer maskout = TMaskImage::New();
      this->SetNthOutput(1, maskout.GetPointer());
    }
    else
    {
      this->SetNumberOfRequiredOutputs(1);
      this->SetNthOutput(1, nullptr);
    }
  }
}

template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
void
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::AllocateOutputs()
{
  if (this->m_GenerateOutputMask)
  {
    // Allocate the output image.
    typename TOutputImage::Pointer output = this->GetOutput();
    output->SetBufferedRegion(output->GetRequestedRegion());
    output->Allocate();
    // Allocate the output mask image.
    typename TMaskImage::Pointer mask = this->GetOutputMask();
    mask->SetBufferedRegion(mask->GetRequestedRegion());
    mask->Allocate();
  }
  else
  {
    Superclass::AllocateOutputs();
  }
}

template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
DataObject::Pointer
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::MakeOutput(
  DataObjectPointerArraySizeType idx)
{
  DataObject::Pointer output;

  switch (idx)
  {
    case 0:
      output = (TOutputImage::New()).GetPointer();
      break;
    case 1:
      output = (TMaskImage::New()).GetPointer();
      break;
  }
  return output.GetPointer();
}

template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
auto
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::GetOutputMask()
  -> MaskImageType *
{
  typename MaskImageType::Pointer res = dynamic_cast<MaskImageType *>(this->ProcessObject::GetOutput(1));
  return res;
}

// a modified version that uses line iterators and only moves the
// histogram in one direction. Hopefully it will be a bit simpler and
// faster due to improved memory access and a tighter loop.
template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
void
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::
  DynamicThreadedGenerateData(const OutputImageRegionType & outputRegionForThread)
{
  HistogramType histogram;
  this->ConfigureHistogram(histogram);

  OutputImageType *      outputImage = this->GetOutput();
  MaskImageType *        outputMask = this->GetOutputMask();
  const InputImageType * inputImage = this->GetInput();
  const MaskImageType *  maskImage = this->GetMaskImage();

  RegionType inputRegion = inputImage->GetRequestedRegion();

  TotalProgressReporter progress(this, inputRegion.GetNumberOfPixels());

  // initialize the histogram
  for (auto listIt = this->m_KernelOffsets.begin(); listIt != this->m_KernelOffsets.end(); ++listIt)
  {
    IndexType idx = outputRegionForThread.GetIndex() + (*listIt);
    if (inputRegion.IsInside(idx) && maskImage->GetPixel(idx) == m_MaskValue)
    {
      histogram.AddPixel(inputImage->GetPixel(idx));
    }
    else
    {
      histogram.AddBoundary();
    }
  }

  // now move the histogram
  FixedArray<short, ImageDimension> direction;
  direction.Fill(1);
  int        axis = ImageDimension - 1;
  OffsetType offset;
  offset.Fill(0);
  RegionType stRegion;
  stRegion.SetSize(this->m_Kernel.GetSize());
  stRegion.PadByRadius(1); // must pad the region by one because of the
                           // translation

  OffsetType centerOffset;
  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    centerOffset[i] = stRegion.GetSize()[i] / 2;
  }

  int BestDirection = this->m_Axes[axis];
  int LineLength = inputRegion.GetSize()[BestDirection];

  // init the offset and get the lists for the best axis
  offset[BestDirection] = direction[BestDirection];
  // it's very important for performances to get a pointer and not a copy
  const OffsetListType * addedList = &this->m_AddedOffsets[offset];
  const OffsetListType * removedList = &this->m_RemovedOffsets[offset];

  using InputLineIteratorType = ImageLinearConstIteratorWithIndex<InputImageType>;
  InputLineIteratorType InLineIt(inputImage, outputRegionForThread);
  InLineIt.SetDirection(BestDirection);
  InLineIt.GoToBegin();
  IndexType LineStart;

  using HistVecType = typename std::vector<HistogramType>;
  HistVecType HistVec(ImageDimension);
  using IndexVecType = typename std::vector<IndexType>;
  IndexVecType PrevLineStartVec(ImageDimension);

  // Steps is used to keep track of the order in which the line
  // iterator passes over the various dimensions.
  int Steps[ImageDimension];

  for (unsigned int i = 0; i < ImageDimension; ++i)
  {
    HistVec[i] = histogram;
    PrevLineStartVec[i] = InLineIt.GetIndex();
    Steps[i] = 0;
  }

  while (!InLineIt.IsAtEnd())
  {
    HistogramType & histRef = HistVec[BestDirection];
    IndexType       PrevLineStart = InLineIt.GetIndex();
    for (InLineIt.GoToBeginOfLine(); !InLineIt.IsAtEndOfLine(); ++InLineIt)
    {
      // Update the histogram
      IndexType currentIdx = InLineIt.GetIndex();

      if (maskImage->GetPixel(currentIdx) == m_MaskValue && histRef.IsValid())
      {
        outputImage->SetPixel(currentIdx,
                              static_cast<OutputPixelType>(histRef.GetValue(inputImage->GetPixel(currentIdx))));
        if (this->m_GenerateOutputMask)
        {
          outputMask->SetPixel(currentIdx, m_MaskValue);
        }
      }
      else
      {
        outputImage->SetPixel(currentIdx, m_FillValue);
        if (this->m_GenerateOutputMask)
        {
          outputMask->SetPixel(currentIdx, m_BackgroundMaskValue);
        }
      }
      stRegion.SetIndex(currentIdx - centerOffset);
      pushHistogram(histRef, addedList, removedList, inputRegion, stRegion, inputImage, maskImage, currentIdx);
    }
    Steps[BestDirection] += LineLength;
    InLineIt.NextLine();
    progress.Completed(outputRegionForThread.GetSize()[0]);
    if (InLineIt.IsAtEnd())
    {
      break;
    }
    LineStart = InLineIt.GetIndex();
    // This section updates the histogram for the next line
    // Since we aren't zig zagging we need to figure out which
    // histogram to update and the direction in which to push
    // it. Then we need to copy that histogram to the relevant
    // places
    OffsetType LineOffset, Changes;
    // Figure out which stored histogram to move and in
    // which direction
    int LineDirection = 0;
    // This function deals with changing planes etc
    this->GetDirAndOffset(LineStart, PrevLineStart, LineOffset, Changes, LineDirection);
    ++(Steps[LineDirection]);
    IndexType              PrevLineStartHist = LineStart - LineOffset;
    const OffsetListType * addedListLine = &this->m_AddedOffsets[LineOffset];
    const OffsetListType * removedListLine = &this->m_RemovedOffsets[LineOffset];
    HistogramType &        tmpHist = HistVec[LineDirection];
    stRegion.SetIndex(PrevLineStart - centerOffset);
    // Now move the histogram
    pushHistogram(
      tmpHist, addedListLine, removedListLine, inputRegion, stRegion, inputImage, maskImage, PrevLineStartHist);

    // PrevLineStartVec[LineDirection] = LineStart;
    // copy the updated histogram and line start entries to the
    // relevant directions. When updating direction 2, for example,
    // new copies of directions 0 and 1 should be made.
    for (unsigned int i = 0; i < ImageDimension; ++i)
    {
      if (Steps[i] > Steps[LineDirection])
      {
        // make sure this is the right thing to do
        HistVec[i] = HistVec[LineDirection];
      }
    }
  }
}

template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>
void
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::pushHistogram(
  HistogramType &        histogram,
  const OffsetListType * addedList,
  const OffsetListType * removedList,
  const RegionType &     inputRegion,
  const RegionType &     kernRegion,
  const InputImageType * inputImage,
  const MaskImageType *  maskImage,
  const IndexType        currentIdx)
{
  if (inputRegion.IsInside(kernRegion))
  {
    // update the histogram
    for (auto addedIt = addedList->begin(); addedIt != addedList->end(); ++addedIt)
    {
      typename InputImageType::IndexType idx = currentIdx + (*addedIt);
      if (maskImage->GetPixel(idx) == m_MaskValue)
      {
        histogram.AddPixel(inputImage->GetPixel(idx));
      }
      else
      {
        histogram.AddBoundary();
      }
    }
    for (auto removedIt = removedList->begin(); removedIt != removedList->end(); ++removedIt)
    {
      typename InputImageType::IndexType idx = currentIdx + (*removedIt);
      if (maskImage->GetPixel(idx) == m_MaskValue)
      {
        histogram.RemovePixel(inputImage->GetPixel(idx));
      }
      else
      {
        histogram.RemoveBoundary();
      }
    }
  }
  else
  {
    // update the histogram
    for (auto addedIt = addedList->begin(); addedIt != addedList->end(); ++addedIt)
    {
      IndexType idx = currentIdx + (*addedIt);
      if (inputRegion.IsInside(idx) && maskImage->GetPixel(idx) == m_MaskValue)
      {
        histogram.AddPixel(inputImage->GetPixel(idx));
      }
      else
      {
        histogram.AddBoundary();
      }
    }
    for (auto removedIt = removedList->begin(); removedIt != removedList->end(); ++removedIt)
    {
      IndexType idx = currentIdx + (*removedIt);
      if (inputRegion.IsInside(idx) && maskImage->GetPixel(idx) == m_MaskValue)
      {
        histogram.RemovePixel(inputImage->GetPixel(idx));
      }
      else
      {
        histogram.RemoveBoundary();
      }
    }
  }
}

template <typename TInputImage, typename TMaskImage, typename TOutputImage, typename TKernel, typename THistogram>

void
MaskedMovingHistogramImageFilter<TInputImage, TMaskImage, TOutputImage, TKernel, THistogram>::PrintSelf(
  std::ostream & os,
  Indent         indent) const
{
  Superclass::PrintSelf(os, indent);

  os << indent << "GenerateOutputMask: " << m_GenerateOutputMask << std::endl;
  os << indent << "FillValue: " << static_cast<typename NumericTraits<OutputPixelType>::PrintType>(m_FillValue)
     << std::endl;
  os << indent << "MaskValue: " << static_cast<typename NumericTraits<MaskPixelType>::PrintType>(m_MaskValue)
     << std::endl;
  os << indent
     << "BackgroundMaskValue: " << static_cast<typename NumericTraits<MaskPixelType>::PrintType>(m_BackgroundMaskValue)
     << std::endl;
}
} // end namespace itk
#endif
