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


namespace itk
{
template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::FastMarchingExtensionImageFilterBase()
{
  m_AuxiliaryAliveValues = nullptr;
  m_AuxiliaryTrialValues = nullptr;

  this->ProcessObject::SetNumberOfRequiredOutputs(1 + AuxDimension);

  AuxImagePointer ptr;
  for (unsigned int k = 0; k < AuxDimension; ++k)
  {
    ptr = AuxImageType::New();
    this->ProcessObject::SetNthOutput(k + 1, ptr.GetPointer());
    this->m_AuxImages[k] = nullptr;
  }
}

template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
void
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::PrintSelf(std::ostream & os,
                                                                                           Indent         indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "Aux alive values: ";
  // os << m_AuxiliaryAliveValues.GetPointer() << std::endl;
  os << indent << "Aux trail values: ";
  // os << m_AuxiliaryTrialValues.GetPointer() << std::endl;
}


template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
auto
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::GetAuxiliaryImage(
  const unsigned int idx) -> AuxImageType *
{
  if (idx >= AuxDimension || this->GetNumberOfIndexedOutputs() < idx + 2)
  {
    return nullptr;
  }

  return static_cast<AuxImageType *>(this->ProcessObject::GetOutput(idx + 1));
}

/*
 *
 */
template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
void
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::GenerateOutputInformation()
{
  // call the superclass implementation of this function
  this->Superclass::GenerateOutputInformation();

  // set the size of all the auxiliary outputs
  // to be the same as the primary output
  OutputImageType * primaryOutput = this->GetOutput();
  for (unsigned int k = 0; k < AuxDimension; ++k)
  {
    AuxImageType * ptr = this->GetAuxiliaryImage(k);
    ptr->CopyInformation(primaryOutput);
  }
}

/*
 *
 */
template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
void
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::EnlargeOutputRequestedRegion(
  DataObject * itkNotUsed(output))
{
  // This filter requires all of the output images in the buffer.
  for (unsigned int j = 0; j < this->GetNumberOfIndexedOutputs(); ++j)
  {
    if (this->ProcessObject::GetOutput(j))
    {
      this->ProcessObject::GetOutput(j)->SetRequestedRegionToLargestPossibleRegion();
    }
  }
}

/*
 *
 */
template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
void
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::InitializeOutput(
  OutputImageType * oImage)
{
  this->Superclass::InitializeOutput(oImage);

  if (!m_AuxiliaryAliveValues)
  {
    itkExceptionMacro("in Initialize(): Null pointer for AuxAliveValues");
  }
  if (m_AuxiliaryAliveValues->size() != (this->m_AlivePoints->size()))
  {
    itkExceptionMacro("in Initialize(): AuxAliveValues is the wrong size");
  }

  if (!m_AuxiliaryTrialValues)
  {
    itkExceptionMacro("in Initialize(): Null pointer for AuxTrialValues");
  }
  if (m_AuxiliaryTrialValues->size() != this->m_TrialPoints->size())
  {
    itkExceptionMacro("in Initialize(): AuxTrialValues is the wrong size");
  }

  // allocate memory for the auxiliary outputs
  for (unsigned int k = 0; k < AuxDimension; ++k)
  {
    AuxImageType * ptr = this->GetAuxiliaryImage(k);
    ptr->SetBufferedRegion(ptr->GetRequestedRegion());
    ptr->Allocate();
    this->m_AuxImages[k] = ptr;
  }

  // set all alive points to alive
  typename Superclass::NodeType node;

  AuxValueVectorType auxVec;

  if (m_AuxiliaryAliveValues)
  {
    AuxValueContainerConstIterator auxIter = m_AuxiliaryAliveValues->Begin();
    NodePairContainerConstIterator pointsIter = this->m_AlivePoints->Begin();
    NodePairContainerConstIterator pointsEnd = this->m_AlivePoints->Begin();

    while (pointsIter != pointsEnd)
    {
      node = pointsIter->Value().GetNode();
      auxVec = auxIter->Value();

      // check if node index is within the output level set
      if (this->m_BufferedRegion.IsInside(node))
      {
        for (unsigned int k = 0; k < AuxDimension; ++k)
        {
          this->m_AuxImages[k]->SetPixel(node, auxVec[k]);
        }
      }
      ++pointsIter;
      ++auxIter;
    } // end container while
  }   // if AuxAliveValues set

  if (m_AuxiliaryTrialValues)
  {
    AuxValueContainerConstIterator auxIter = m_AuxiliaryTrialValues->Begin();
    NodePairContainerConstIterator pointsIter = this->m_TrialPoints->Begin();
    NodePairContainerConstIterator pointsEnd = this->m_TrialPoints->End();

    while (pointsIter != pointsEnd)
    {
      node = pointsIter->Value().GetNode();
      auxVec = auxIter->Value();

      // check if node index is within the output level set
      if (this->m_BufferedRegion.IsInside(node))
      {
        for (unsigned int k = 0; k < AuxDimension; ++k)
        {
          this->m_AuxImages[k]->SetPixel(node, auxVec[k]);
        }
      }
      ++pointsIter;
      ++auxIter;
    } // end container loop
  }   // if AuxTrialValues set
}

template <typename TInput, typename TOutput, typename TAuxValue, unsigned int VAuxDimension>
void
FastMarchingExtensionImageFilterBase<TInput, TOutput, TAuxValue, VAuxDimension>::UpdateValue(OutputImageType * oImage,
                                                                                             const NodeType &  iNode)
{
  // A extension value at node is chosen such that
  // grad(F) dot_product grad(Phi) = 0
  // where F is the extended speed function and Phi is
  // the level set function.
  //
  // The extension value can approximated as a weighted
  // sum of the values from nodes used in the calculation
  // of the distance by the superclass.
  //
  // For more detail see Chapter 11 of
  // "Level Set Methods and Fast Marching Methods", J.A. Sethian,
  // Cambridge Press, Second edition, 1999.

  typename Superclass::InternalNodeStructureArray NodesUsed;

  this->GetInternalNodesUsed(oImage, iNode, NodesUsed);

  InternalNodeStructure temp_node;

  auto outputPixel = static_cast<OutputPixelType>(this->Solve(oImage, iNode, NodesUsed));

  if (outputPixel < this->m_LargeValue)
  {
    oImage->SetPixel(iNode, outputPixel);

    // insert point into trial heap
    this->m_LabelImage->SetPixel(iNode, Traits::Trial);

    // node.SetValue( outputPixel );
    // node.SetIndex( index );
    // m_TrialHeap.push(node);
    this->m_Heap.push(NodePairType(iNode, outputPixel));

    // update auxiliary values
    for (unsigned int k = 0; k < AuxDimension; ++k)
    {
      double       numer = 0.;
      double       denom = 0.;
      AuxValueType auxVal;

      for (unsigned int j = 0; j < ImageDimension; ++j)
      {
        temp_node = NodesUsed[j];

        if (outputPixel < temp_node.m_Value)
        {
          break;
        }

        auxVal = this->m_AuxImages[k]->GetPixel(temp_node.m_Node);
        numer += auxVal * (outputPixel - temp_node.m_Value);
        denom += outputPixel - temp_node.m_Value;
      }

      if (denom > itk::Math::eps)
      {
        auxVal = static_cast<AuxValueType>(numer / denom);
      }
      else
      {
        auxVal = AuxValueType{};
      }

      this->m_AuxImages[k]->SetPixel(iNode, auxVal);
    }
  }
}
} // namespace itk

#endif
