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

#include "itkInPlaceImageFilter.h"
#include "itkFiniteDifferenceFunction.h"
#include "vnl/vnl_vector.h"
#include "itkImageRegionIterator.h"

#include "itkListSample.h"
#include "itkKdTreeGenerator.h"

namespace itk
{
/**
 * \class MultiphaseFiniteDifferenceImageFilter
 *
 * \par The Finite Difference Solver Hierarchy
 *
 * This is an alternate version of the ITK finite difference solver (FDS)
 * framework, supporting the solution of multiple functions, simultaneously.
 * The FDS framework is a set of classes for creating filters to solve partial
 * differential equations on images using an iterative, finite difference
 * update scheme.
 *
 * \par
 * The high-level algorithm implemented by the framework can be described by
 * the following pseudocode.
 *
   \code
    WHILE NOT convergence:
       FOR ALL pixels i
        FOR ALL functions f
        min_time_step = min(min_time_step, calculate_change(f, i))
        FOR ALL functions f
            update(f, i, time_step)
   \endcode
 *
 * \par
 * The following equation describes update \f$n+1\f$ at pixel \f$i\f$ on
 * discrete image \f$ u \f$ :
 *
 * \par
 * \f$u_{\mathbf{i}}^{n+1}=u^n_{\mathbf{i}}+\Delta u^n_{\mathbf{i}}\Delta t\f$
 *
 * \par Component objects
 * The FDS hierarchy is comprised of two component object types, variations of
 * which are designed to be plugged together to create filters for different
 * applications.  At the process level are the "solver" objects, which are
 * subclasses of MultiphaseFiniteDifferenceImageFilter.  Solver objects are filters that
 * take image inputs and produce image outputs.  Solver objects require a
 * "finite difference function" object to perform the calculation at each
 * image pixel during iteration.  These specialized function objects are
 * subclasses of FiniteDifferenceFunction. FiniteDifferenceFunctions take a
 * neighborhood of pixels as input (in the form of an
 * itk::NeighborhoodIterator) and produce a scalar valued result.
 *
 * \par
 * Filters for different applications are created by defining a function object
 * to handle the numerical calculations and choosing (or creating) a solver
 * object that reflects the requirements and constraints of the application.
 * For example, anisotropic diffusion filters are created by plugging
 * anisotropic diffusion functions into the DenseFiniteDifferenceImageFilter2.
 * The separation between function object and solver object allows us to
 * create, for example, sparse-field, dense-field, and narrow-band
 * implementations of a level-set surface evolution filter can all be
 * constructed by plugging the same function object into three different,
 * specialized solvers.
 *
 * \par Creating new filters in this hierarchy
 * The procedure for creating a filter within the FDS hierarchy is to identify
 * all the virtual methods that need to be defined for your particular
 * application.  In the simplest case, a filter needs only to instantiate a
 * specific function object and define some halting criteria.  For more
 * complicated applications, you may need to define a specialized type of
 * iteration scheme or updating procedure in a higher-level solver object.
 *
 * \par
 * Some simple examples are the specific subclasses of
 * AnisotropicDiffusionImageFilter.  The leaves of the anisotropic diffusion
 * filter tree only define the function object they use for their particular
 * flavor of diffusion.  See CurvatureAnisotropicDiffusionImageFilter and
 * GradientAnisotropicDiffusionImageFilter for details.
 *
 * \par FiniteDifferenceImageFilter2
 * This class defines the generic solver API at the top level of the FDS
 * framework. FiniteDifferenceImageFilter2 is an abstract class that implements
 * the generic, high-level algorithm (described above).
 *
 * \par Inputs and Outputs
 * This filter is an Image to Image filter.  Depending on the specific
 * subclass implementation, finite difference image filters may process a
 * variety of image types.  The input to the filter is the initial
 * value of \f$ u \f$ and the output of the filter is the solution to the
 * p.d.e.
 *
 * \par How to use this class
 * GenerateData() relies on several virtual methods that must be defined by a
 * subclass.  Specifically: \em AllocateUpdateBuffer \em ApplyUpdate
 * \em CalculateChange and \em Halt.  To create a finite difference solver,
 * implement a subclass to define these methods.
 *
 * \par
 * Note that there is no fixed container type for the buffer used to hold the
 * update \f$ \Delta \f$.  The container might be another image, or simply a
 * list of values.  AllocateUpdateBuffer is responsible for creating the \f$
 * \Delta \f$ container.  CalculateChange populates this buffer and ApplyUpdate
 * adds the buffer values to the output image (solution).  The boolean Halt()
 * (or ThreadedHalt) method returns a true value to stop iteration.
 *
 *
 * Based on the paper:
 *
 *        "An active contour model without edges"
 *         T. Chan and L. Vese.
 *         In Scale-Space Theories in Computer Vision, pages 141-151, 1999.
 *
 * \author Mosaliganti K., Smith B., Gelas A., Gouaillard A., Megason S.
 *
 *  This code was taken from the Insight Journal paper:
 *
 *      "Cell Tracking using Coupled Active Surfaces for Nuclei and Membranes"
 *      http://www.insight-journal.org/browse/publication/642
 *      https://hdl.handle.net/10380/3055
 *
 *  That is based on the papers:
 *
 *      "Level Set Segmentation: Active Contours without edge"
 *      http://www.insight-journal.org/browse/publication/322
 *      https://hdl.handle.net/1926/1532
 *
 *      and
 *
 *      "Level set segmentation using coupled active surfaces"
 *      http://www.insight-journal.org/browse/publication/323
 *      https://hdl.handle.net/1926/1533
 *
 *
 * \ingroup ImageFilter
 * \ingroup LevelSetSegmentation
 * \sa DenseFiniteDifferenceImageFilter2
 * \ingroup ITKReview
 */
template <typename TInputImage,
          typename TFeatureImage,
          typename TOutputImage,
          typename TFiniteDifferenceFunction = FiniteDifferenceFunction<TOutputImage>,
          typename TIdCell = unsigned int>
class ITK_TEMPLATE_EXPORT MultiphaseFiniteDifferenceImageFilter : public InPlaceImageFilter<TFeatureImage, TOutputImage>
{
public:
  ITK_DISALLOW_COPY_AND_ASSIGN(MultiphaseFiniteDifferenceImageFilter);

  /** Standard class type aliases. */
  using Self = MultiphaseFiniteDifferenceImageFilter;
  using Superclass = InPlaceImageFilter<TFeatureImage, TOutputImage>;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;

  /** Run-time type information (and related methods) */
  itkTypeMacro(MultiphaseFiniteDifferenceImageFilter, InPlaceImageFilter);

  /** Dimensionality of input and output data is assumed to be the same. */
  static constexpr unsigned int ImageDimension = TOutputImage::ImageDimension;

  /** Input and output image types. */
  using InputImageType = TInputImage;
  using InputImagePointer = typename InputImageType::Pointer;
  using InputPointType = typename InputImageType::PointType;
  using InputCoordRepType = typename InputPointType::CoordRepType;
  using InputIndexType = typename InputImageType::IndexType;
  using InputIndexValueType = typename InputIndexType::IndexValueType;
  using InputSizeType = typename InputImageType::SizeType;
  using InputSizeValueType = typename InputSizeType::SizeValueType;
  using InputRegionType = typename InputImageType::RegionType;
  using InputPixelType = typename InputImageType::PixelType;
  using InputSpacingType = typename InputImageType::SpacingType;
  using InputOffsetValueType = typename InputImageType::OffsetValueType;

  using FeatureImageType = TFeatureImage;
  using FeatureImagePointer = typename FeatureImageType::Pointer;
  using FeatureRegionType = typename FeatureImageType::RegionType;
  using FeatureSizeType = typename FeatureImageType::SizeType;
  using FeatureSpacingType = typename FeatureImageType::SpacingType;
  using FeaturePointType = typename FeatureImageType::PointType;
  using FeaturePixelType = typename FeatureImageType::PixelType;

  using OutputImageType = TOutputImage;
  using OutputImagePointer = typename OutputImageType::Pointer;
  using OutputPixelType = typename OutputImageType::PixelType;
  using OutputRegionType = typename OutputImageType::RegionType;
  using OutputSizeType = typename OutputImageType::SizeType;
  using OutputSizeValueType = typename OutputImageType::SizeValueType;
  using OutputIndexType = typename OutputImageType::IndexType;
  using OutputIndexValueType = typename OutputImageType::IndexValueType;

  using IdCellType = TIdCell;
  using VectorIdCellType = std::vector<IdCellType>;

  /** The value type of the time step.  This is distinct from PixelType
   * because PixelType may often be a vector value, while the TimeStep is
   * a scalar value. */
  using FiniteDifferenceFunctionType = TFiniteDifferenceFunction;
  using FiniteDifferenceFunctionPointer = typename FiniteDifferenceFunctionType::Pointer;
  using TimeStepType = typename FiniteDifferenceFunctionType::TimeStepType;
  using TimeStepVectorType = typename std::vector<TimeStepType>;
  using RadiusType = typename FiniteDifferenceFunctionType::RadiusType;

  using CentroidVectorType = Vector<float, Self::ImageDimension>;
  using SampleType = Statistics::ListSample<CentroidVectorType>;
  using KdTreeGeneratorType = Statistics::KdTreeGenerator<SampleType>;
  using KdTreeGeneratorPointer = typename KdTreeGeneratorType::Pointer;
  using KdTreeType = typename KdTreeGeneratorType::KdTreeType;
  using KdTreePointer = typename KdTreeType::Pointer;

  /** This method returns a pointer to a FiniteDifferenceFunction object that
   * will be used by the filter to calculate updates at image pixels.
   * \param functionIndex Index of difference function to return.
   * \returns A FiniteDifferenceObject pointer. */
  virtual const FiniteDifferenceFunctionPointer
  GetDifferenceFunction(const IdCellType & functionIndex) const
  {
    if (functionIndex < m_FunctionCount)
    {
      return (this->m_DifferenceFunctions[functionIndex]);
    }
    else
    {
      return nullptr;
    }
  }

  /** This method sets the pointer to a FiniteDifferenceFunction object that
   * will be used by the filter to calculate updates at image pixels.
   * \param functionIndex Index of difference function to set.
   * \param function Pointer to difference function to set. */
  virtual void
  SetDifferenceFunction(const IdCellType & functionIndex, FiniteDifferenceFunctionPointer function)
  {
    if (functionIndex < m_FunctionCount)
    {
      this->m_DifferenceFunctions[functionIndex] = function;
    }
  }

  /** Set/Get the number of iterations that the filter will run. */
  itkSetMacro(NumberOfIterations, unsigned int);
  itkGetConstReferenceMacro(NumberOfIterations, unsigned int);

  /** Use the image spacing information in calculations. Use this option if you
   *  want derivatives in physical space. Default is UseImageSpacingOn. */
  itkSetMacro(UseImageSpacing, bool);
  itkBooleanMacro(UseImageSpacing);
  itkGetConstReferenceMacro(UseImageSpacing, bool);

  /** Set/Get the maximum error allowed in the solution.  This may not be
   * defined for all solvers and its meaning may change with the application. */
  itkSetMacro(MaximumRMSError, double);
  itkGetConstReferenceMacro(MaximumRMSError, double);

  /** Set/Get the root mean squared change of the previous iteration. May not
      be used by all solvers. */
  itkSetMacro(RMSChange, double);
  itkGetConstReferenceMacro(RMSChange, double);

  /** Set/Get the state of the filter. */
  itkSetMacro(InitializedState, bool);
  itkGetConstReferenceMacro(InitializedState, bool);
  itkBooleanMacro(InitializedState);

  /** Require the filter to be manually reinitialized (by calling
      SetInitializedStateOff() */
  itkSetMacro(ManualReinitialization, bool);
  itkGetConstReferenceMacro(ManualReinitialization, bool);
  itkBooleanMacro(ManualReinitialization);

  /** Set the number of elapsed iterations of the filter. */
  itkSetMacro(ElapsedIterations, unsigned int);

  /** Get the number of elapsed iterations of the filter. */
  itkGetConstReferenceMacro(ElapsedIterations, unsigned int);

  void
  SetLevelSet(const IdCellType & i, const InputImageType * levelSet)
  {
    m_LevelSet[i] = InputImageType::New();
    m_LevelSet[i]->SetRequestedRegion(levelSet->GetRequestedRegion());
    m_LevelSet[i]->SetBufferedRegion(levelSet->GetBufferedRegion());
    m_LevelSet[i]->SetLargestPossibleRegion(levelSet->GetLargestPossibleRegion());
    m_LevelSet[i]->Allocate();
    m_LevelSet[i]->CopyInformation(levelSet);

    ImageRegionConstIterator<InputImageType> in(levelSet, levelSet->GetBufferedRegion());
    ImageRegionIterator<InputImageType>      cp(m_LevelSet[i], levelSet->GetBufferedRegion());

    in.GoToBegin();
    cp.GoToBegin();

    while (!in.IsAtEnd())
    {
      cp.Set(in.Get());
      ++in;
      ++cp;
    }
  }

  InputImagePointer
  GetLevelSet(const IdCellType & i)
  {
    if (i >= m_FunctionCount)
    {
      itkExceptionMacro("Request for level set #" << i << " but there are only " << m_FunctionCount);
    }
    else
    {
      return m_LevelSet[i];
    }
  }

  void
  SetLookup(VectorIdCellType lookup)
  {
    this->m_Lookup = lookup;
  }

  void
  SetKdTree(KdTreeType * kdtree)
  {
    this->m_KdTree = kdtree;
  }

  void
  SetFunctionCount(const IdCellType & n)
  {
    m_FunctionCount = n;

    m_DifferenceFunctions.resize(m_FunctionCount, nullptr);

    RadiusType radius;
    radius.Fill(1);

    for (unsigned int i = 0; i < this->m_FunctionCount; i++)
    {
      this->m_DifferenceFunctions[i] = FiniteDifferenceFunctionType::New();
      this->m_DifferenceFunctions[i]->Initialize(radius);
    }

    // Initialize the images
    m_LevelSet.resize(m_FunctionCount, nullptr);

    // Initialize the lookup table
    this->m_Lookup.resize(m_FunctionCount);

    IdCellType k = 1;

    {
      auto it = this->m_Lookup.begin();
      auto _end = this->m_Lookup.end();
      while (it != _end)
      {
        *it = k;
        ++it;
        ++k;
      }
    }
  }

protected:
  MultiphaseFiniteDifferenceImageFilter()
  {
    this->m_KdTree = nullptr;
    this->m_ElapsedIterations = 0;
    this->m_MaximumRMSError = itk::Math::eps;
    this->m_RMSChange = NumericTraits<double>::max();
    this->m_UseImageSpacing = true;
    this->m_ManualReinitialization = false;
    this->m_InitializedState = false;
    this->m_NumberOfIterations = NumericTraits<unsigned int>::max();
    this->m_FunctionCount = 0;
    this->InPlaceOff();
  }

  ~MultiphaseFiniteDifferenceImageFilter() override = default;

  IdCellType                     m_FunctionCount;
  std::vector<InputImagePointer> m_LevelSet;
  VectorIdCellType               m_Lookup;
  KdTreePointer                  m_KdTree;

  unsigned int m_ElapsedIterations;
  double       m_MaximumRMSError;
  double       m_RMSChange;
  unsigned int m_NumberOfIterations;

  /** The function that will be used in calculating updates for each pixel. */
  std::vector<FiniteDifferenceFunctionPointer> m_DifferenceFunctions;

  /** Control whether derivatives use spacing of the input image in its
   * calculation. */
  bool m_UseImageSpacing;

  void
  PrintSelf(std::ostream & os, Indent indent) const override;

  /** This method allocates a temporary update container in the subclass. */
  virtual void
  AllocateUpdateBuffer() = 0;

  /** This method is defined by a subclass to apply changes to the output
   * from an update buffer and a time step value "dt".
   * \param dt Time step value. */
  virtual void
  ApplyUpdate(TimeStepType dt) = 0;

  /** This method is defined by a subclass to populate an update buffer
   * with changes for the pixels in the output.  It returns a time
   * step value to be used for the update.
   * \returns A time step to use in updating the output with the changes
   * calculated from this method. */
  virtual TimeStepType
  CalculateChange() = 0;

  /** This method can be defined in subclasses as needed to copy the input
   * to the output. See DenseFiniteDifferenceImageFilter2 for an
   * implementation. */
  virtual void
  CopyInputToOutput() = 0;

  /** This is the default, high-level algorithm for calculating finite
   * difference solutions.  It calls virtual methods in its subclasses
   * to implement the major steps of the algorithm. */
  void
  GenerateData() override;

  /** FiniteDifferenceImageFilter2 needs a larger input requested region than
   * the output requested region.  As such, we need to provide
   * an implementation for GenerateInputRequestedRegion() in order to inform
   * the pipeline execution model.
   *
   * \par
   * The filter will ask for a padded region to perform its neighborhood
   * calculations.  If no such region is available, the boundaries will be
   * handled as described in the FiniteDifferenceFunction defined by the
   * subclass.
   * \sa ProcessObject::GenerateInputRequestedRegion() */
  void
  GenerateInputRequestedRegion() override;

  /** This method returns true when the current iterative solution of the
   * equation has met the criteria to stop solving.  Defined by a subclass. */
  virtual bool
  Halt();

  /** This method is similar to Halt(), and its default implementation in this
   * class is simply to call Halt(). However, this method takes as a parameter a
   * void pointer to the MultiThreaderBase::WorkUnitInfo structure. If you
   * override this method instead of overriding Halt, you will be able to get the
   * current thread ID and handle the Halt method accordingly. This is useful if
   * you are doing a lot of processing in Halt that you don't want parallelized.
   * Notice that ThreadedHalt is only called by the multithreaded filters, so you
   * still should implement Halt, just in case a non-threaded filter is used.
   */
  virtual bool
  ThreadedHalt(void * itkNotUsed(threadInfo))
  {
    return this->Halt();
  }

  /** This method is optionally defined by a subclass and is called before
   * the loop of iterations of calculate_change & update. It does the global
   * initialization, i.e. in the SparseFieldLevelSetImageFilter, initialize
   * the list of layers.
   * */
  virtual void
  Initialize()
  {}

  /** This method is optionally defined by a subclass and is called immediately
   * prior to each iterative CalculateChange-ApplyUpdate cycle.  It can be
   * used to set global variables needed for the next iteration (ie. average
   * gradient magnitude of the image in anisotropic diffusion functions), or
   * otherwise prepare for the next iteration.
   * */
  virtual void
  InitializeIteration()
  {
    for (IdCellType i = 0; i < this->m_FunctionCount; i++)
    {
      this->m_DifferenceFunctions[i]->InitializeIteration();
    }
  }

  /** Virtual method for resolving a single time step from a set of time steps
   * returned from processing threads.
   * \return Time step (dt) for the iteration update based on a list
   * of time steps generated from the threaded calculated change method (one
   * for each region processed).
   *
   * \param timeStepList The set of time changes compiled from all the threaded
   *        calls to ThreadedGenerateData.
   *
   * \param valid The set of flags indicating which of "list" elements are valid
   *
   * The default is to return the minimum value in the list. */
  inline TimeStepType
  ResolveTimeStep(const TimeStepVectorType & timeStepList, const std::vector<bool> & valid);

  /** This method is called after the solution has been generated to allow
   * subclasses to apply some further processing to the output. */
  virtual void
  PostProcessOutput()
  {}

private:
  /** Indicates whether the filter automatically resets to UNINITIALIZED state
      after completing, or whether filter must be manually reset */
  bool m_ManualReinitialization;

  /** State that the filter is in, i.e. UNINITIALIZED or INITIALIZED */
  bool m_InitializedState;
};
} // end namespace itk

#ifndef ITK_MANUAL_INSTANTIATION
#  include "itkMultiphaseFiniteDifferenceImageFilter.hxx"
#endif

#endif
