// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2011-2014 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#ifndef EIGEN_BASIC_PRECONDITIONERS_H
#define EIGEN_BASIC_PRECONDITIONERS_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

namespace Eigen {

/** \ingroup IterativeLinearSolvers_Module
  * \brief A preconditioner based on the diagonal entries
  *
  * This class allows to approximately solve for A.x = b problems assuming A is a diagonal matrix.
  * In other words, this preconditioner neglects all off diagonal entries and, in Eigen's language, solves for:
    \code
    A.diagonal().asDiagonal() . x = b
    \endcode
  *
  * \tparam Scalar_ the type of the scalar.
  *
  * \implsparsesolverconcept
  *
  * This preconditioner is suitable for both selfadjoint and general problems.
  * The diagonal entries are pre-inverted and stored into a dense vector.
  *
  * \note A variant that has yet to be implemented would attempt to preserve the norm of each column.
  *
  * \sa class LeastSquareDiagonalPreconditioner, class ConjugateGradient
  */
template <typename Scalar_>
class DiagonalPreconditioner {
  using Scalar = Scalar_;
  using Vector = Matrix<Scalar, Dynamic, 1>;

 public:
  using StorageIndex = typename Vector::StorageIndex;
  enum { ColsAtCompileTime = Dynamic, MaxColsAtCompileTime = Dynamic };

  DiagonalPreconditioner() = default;

  template <typename MatType>
  explicit DiagonalPreconditioner(const MatType& mat) : m_invdiag(mat.cols()) {
    compute(mat);
  }

  constexpr Index rows() const noexcept { return m_invdiag.size(); }
  constexpr Index cols() const noexcept { return m_invdiag.size(); }

  template <typename MatType>
  DiagonalPreconditioner& analyzePattern(const MatType&) {
    return *this;
  }

  template <typename MatType>
  DiagonalPreconditioner& factorize(const MatType& mat) {
    m_invdiag.resize(mat.cols());
    for (Index j = 0; j < mat.outerSize(); ++j) {
      typename MatType::InnerIterator it(mat, j);
      while (it && it.index() != j) ++it;
      if (it && it.index() == j && it.value() != Scalar(0))
        m_invdiag(j) = Scalar(1) / it.value();
      else
        m_invdiag(j) = Scalar(1);
    }
    m_isInitialized = true;
    return *this;
  }

  template <typename MatType>
  DiagonalPreconditioner& compute(const MatType& mat) {
    return factorize(mat);
  }

  /** \internal */
  template <typename Rhs, typename Dest>
  void _solve_impl(const Rhs& b, Dest& x) const {
    x = m_invdiag.array() * b.array();
  }

  template <typename Rhs>
  inline Solve<DiagonalPreconditioner, Rhs> solve(const MatrixBase<Rhs>& b) const {
    eigen_assert(m_isInitialized && "DiagonalPreconditioner is not initialized.");
    EIGEN_INTERNAL_ASSERT_SIZE_MATCHES_ROWS_MSG(
        m_invdiag, b, "DiagonalPreconditioner::solve(): invalid number of rows of the right hand side matrix b");
    return Solve<DiagonalPreconditioner, Rhs>(*this, b.derived());
  }

  ComputationInfo info() const { return Success; }

 protected:
  Vector m_invdiag;
  bool m_isInitialized = false;
};

/** \ingroup IterativeLinearSolvers_Module
  * \brief Jacobi preconditioner for LeastSquaresConjugateGradient
  *
  * This class allows to approximately solve for A' A x  = A' b problems assuming A' A is a diagonal matrix.
  * In other words, this preconditioner neglects all off diagonal entries and, in Eigen's language, solves for:
    \code
    (A.adjoint() * A).diagonal().asDiagonal() * x = b
    \endcode
  *
  * \tparam Scalar_ the type of the scalar.
  *
  * \implsparsesolverconcept
  *
  * The diagonal entries are pre-inverted and stored into a dense vector.
  *
  * \sa class LeastSquaresConjugateGradient, class DiagonalPreconditioner
  */
template <typename Scalar_>
class LeastSquareDiagonalPreconditioner : public DiagonalPreconditioner<Scalar_> {
  using Scalar = Scalar_;
  using RealScalar = typename NumTraits<Scalar>::Real;
  using Base = DiagonalPreconditioner<Scalar_>;
  using Base::m_invdiag;

 public:
  LeastSquareDiagonalPreconditioner() = default;

  template <typename MatType>
  explicit LeastSquareDiagonalPreconditioner(const MatType& mat) : Base() {
    compute(mat);
  }

  template <typename MatType>
  LeastSquareDiagonalPreconditioner& analyzePattern(const MatType&) {
    return *this;
  }

  template <typename MatType>
  LeastSquareDiagonalPreconditioner& factorize(const MatType& mat) {
    // Compute the inverse squared-norm of each column of mat
    m_invdiag.resize(mat.cols());
    EIGEN_IF_CONSTEXPR (MatType::IsRowMajor) {
      m_invdiag.setZero();
      for (Index j = 0; j < mat.outerSize(); ++j) {
        for (typename MatType::InnerIterator it(mat, j); it; ++it) m_invdiag(it.index()) += numext::abs2(it.value());
      }
      for (Index j = 0; j < mat.cols(); ++j) {
        RealScalar sum = numext::real(m_invdiag(j));
        m_invdiag(j) = sum > RealScalar(0) ? RealScalar(1) / sum : RealScalar(1);
      }
    } else {
      for (Index j = 0; j < mat.outerSize(); ++j) {
        RealScalar sum = mat.col(j).squaredNorm();
        m_invdiag(j) = sum > RealScalar(0) ? RealScalar(1) / sum : RealScalar(1);
      }
    }
    Base::m_isInitialized = true;
    return *this;
  }

  template <typename MatType>
  LeastSquareDiagonalPreconditioner& compute(const MatType& mat) {
    return factorize(mat);
  }

  ComputationInfo info() const { return Success; }
};

/** \ingroup IterativeLinearSolvers_Module
 * \brief A naive preconditioner which approximates any matrix as the identity matrix
 *
 * \implsparsesolverconcept
 *
 * \sa class DiagonalPreconditioner
 */
class IdentityPreconditioner {
 public:
  IdentityPreconditioner() = default;

  template <typename MatrixType>
  explicit IdentityPreconditioner(const MatrixType&) {}

  template <typename MatrixType>
  IdentityPreconditioner& analyzePattern(const MatrixType&) {
    return *this;
  }

  template <typename MatrixType>
  IdentityPreconditioner& factorize(const MatrixType&) {
    return *this;
  }

  template <typename MatrixType>
  IdentityPreconditioner& compute(const MatrixType&) {
    return *this;
  }

  template <typename Rhs>
  inline const Rhs& solve(const Rhs& b) const {
    return b;
  }

  ComputationInfo info() const { return Success; }
};

}  // end namespace Eigen

#endif  // EIGEN_BASIC_PRECONDITIONERS_H