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


namespace itk
{

template <typename TInputMesh, typename TOutputMesh, typename TSolverTraits>
LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints<TInputMesh, TOutputMesh, TSolverTraits>::
  LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints()
  : m_Lambda(1)
  , m_LambdaSquare(1)
{}

template <typename TInputMesh, typename TOutputMesh, typename TSolverTraits>
void
LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints<TInputMesh, TOutputMesh, TSolverTraits>::
  ComputeVertexIdMapping()
{
  OutputMeshType * output = this->GetOutput();

  typename OutputMeshType::PointsContainer * points = output->GetPoints();

  typename OutputMeshType::PointsContainerIterator       pIt = points->Begin();
  const typename OutputMeshType::PointsContainerIterator pEnd = points->End();

  OutputPointIdentifier k = 0;

  while (pIt != pEnd)
  {
    this->m_InternalMap.insert(typename OutputMapPointIdentifier::value_type(pIt->Index(), k++));
    ++pIt;
  }
}

template <typename TInputMesh, typename TOutputMesh, typename TSolverTraits>
void
LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints<TInputMesh, TOutputMesh, TSolverTraits>::SetLocalLambda(
  OutputPointIdentifier vId,
  OutputCoordRepType    iL)
{
  m_LocalLambdaSquare[vId] = iL * iL;
}


template <typename TInputMesh, typename TOutputMesh, typename TSolverTraits>
void
LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints<TInputMesh, TOutputMesh, TSolverTraits>::FillMatrix(
  MatrixType & iM,
  VectorType & iBx,
  VectorType & iBy,
  VectorType & iBz)
{
  OutputMeshType * output = this->GetOutput();

  OutputMapPointIdentifierConstIterator it = this->m_InternalMap.begin();
  OutputMapPointIdentifierConstIterator end = this->m_InternalMap.end();

  while (it != end)
  {
    const OutputPointIdentifier vId1 = it->first;
    const auto                  internalId1 = static_cast<unsigned int>(it->second);

    RowType row;
    this->FillMatrixRow(vId1, this->m_Order, NumericTraits<OutputCoordRepType>::OneValue(), row);

    RowConstIterator rIt = row.begin();
    RowConstIterator rEnd = row.end();

    while (rIt != rEnd)
    {
      const OutputPointIdentifier vId2 = rIt->first;
      const OutputCoordRepType    weight = rIt->second;

      const OutputPointType p = output->GetPoint(vId2);
      iBx[internalId1] += weight * p[0];
      iBy[internalId1] += weight * p[1];
      iBz[internalId1] += weight * p[2];

      const auto internalId2 = static_cast<unsigned int>(this->m_InternalMap[vId2]);
      SolverTraits::FillMatrix(iM, internalId1, internalId2, weight);

      ++rIt;
    }
    ++it;
  }
}

template <typename TInputMesh, typename TOutputMesh, typename TSolverTraits>
void
LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints<TInputMesh, TOutputMesh, TSolverTraits>::GenerateData()
{
  this->m_LambdaSquare = this->m_Lambda * this->m_Lambda;

  this->CopyInputMeshToOutputMesh();

  if (!this->m_Constraints.empty())
  {
    OutputMeshType * output = this->GetOutput();

    this->m_CoefficientMap.clear();
    this->m_MixedAreaMap.clear();

    this->ComputeVertexIdMapping();

    const auto N = static_cast<unsigned int>(this->m_InternalMap.size());

    MatrixType M = SolverTraits::InitializeSparseMatrix(N, N);

    VectorType Bx = SolverTraits::InitializeVector(N);
    Bx.fill(0.);

    VectorType By = SolverTraits::InitializeVector(N);
    By.fill(0.);

    VectorType Bz = SolverTraits::InitializeVector(N);
    Bz.fill(0.);

    this->FillMatrix(M, Bx, By, Bz);

    MatrixType Mt(M.transpose());

    MatrixType A(Mt * M);

    VectorType Cx, Cy, Cz;
    Mt.mult(Bx, Cx);
    Mt.mult(By, Cy);
    Mt.mult(Bz, Cz);

    typename OutputMeshType::PointsContainer * points = output->GetPoints();

    for (ConstraintMapConstIterator cIt = this->m_Constraints.begin(); cIt != this->m_Constraints.end(); ++cIt)
    {
      const OutputPointIdentifier vId = cIt->first;
      OutputPointType             p = points->GetElement(vId);
      p += cIt->second;

      const OutputPointIdentifier index = this->m_InternalMap[vId];

      OutputCoordRepType l2 = m_LambdaSquare;

      typename std::unordered_map<OutputPointIdentifier, OutputCoordRepType>::const_iterator lambdaIt =
        this->m_LocalLambdaSquare.find(vId);
      if (lambdaIt != this->m_LocalLambdaSquare.end())
      {
        l2 = lambdaIt->second;
      }

      Cx[index] += l2 * p[0];
      Cy[index] += l2 * p[1];
      Cz[index] += l2 * p[2];

      SolverTraits::AddToMatrix(A, index, index, l2);
    }

    VectorType X = SolverTraits::InitializeVector(N);
    X.fill(0.);

    VectorType Y = SolverTraits::InitializeVector(N);
    Y.fill(0.);

    VectorType Z = SolverTraits::InitializeVector(N);
    Z.fill(0.);

    this->SolveLinearSystems(A, Cx, Cy, Cz, X, Y, Z);

    OutputMapPointIdentifierConstIterator it = this->m_InternalMap.begin();
    OutputMapPointIdentifierConstIterator end = this->m_InternalMap.end();

    while (it != end)
    {
      const OutputPointIdentifier vId = it->first;
      const auto                  internalId = static_cast<unsigned int>(it->second);

      OutputPointType & p = points->ElementAt(vId);

      auto x = static_cast<OutputCoordRepType>(X[internalId]);
      p[0] = x;

      auto y = static_cast<OutputCoordRepType>(Y[internalId]);
      p[1] = y;

      auto z = static_cast<OutputCoordRepType>(Z[internalId]);
      p[2] = z;

      ++it;
    }
  }
}

template <typename TInputMesh, typename TOutputMesh, typename TSolverTraits>
void
LaplacianDeformationQuadEdgeMeshFilterWithSoftConstraints<TInputMesh, TOutputMesh, TSolverTraits>::PrintSelf(
  std::ostream & os,
  Indent         indent) const
{
  Superclass::PrintSelf(os, indent);
}
} // end namespace itk

#endif
