// Ceres Solver - A fast non-linear least squares minimizer
// Copyright 2023 Google Inc. All rights reserved.
// http://ceres-solver.org/
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice,
//   this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
// * Neither the name of Google Inc. nor the names of its contributors may be
//   used to endorse or promote products derived from this software without
//   specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Author: sameeragarwal@google.com (Sameer Agarwal)
//
// Abstract interface for objects solving linear systems of various
// kinds.

#ifndef CERES_INTERNAL_LINEAR_SOLVER_H_
#define CERES_INTERNAL_LINEAR_SOLVER_H_

#include <cstddef>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include "absl/log/check.h"
#include "ceres/block_sparse_matrix.h"
#include "ceres/casts.h"
#include "ceres/compressed_row_sparse_matrix.h"
#include "ceres/context_impl.h"
#include "ceres/dense_sparse_matrix.h"
#include "ceres/execution_summary.h"
#include "ceres/internal/disable_warnings.h"
#include "ceres/internal/export.h"
#include "ceres/triplet_sparse_matrix.h"
#include "ceres/types.h"

namespace ceres::internal {

enum class LinearSolverTerminationType {
  // Termination criterion was met.
  SUCCESS,

  // Solver ran for max_num_iterations and terminated before the
  // termination tolerance could be satisfied.
  NO_CONVERGENCE,

  // Solver was terminated due to numerical problems, generally due to
  // the linear system being poorly conditioned.
  FAILURE,

  // Solver failed with a fatal error that cannot be recovered from,
  // e.g. CHOLMOD ran out of memory when computing the symbolic or
  // numeric factorization or an underlying library was called with
  // the wrong arguments.
  FATAL_ERROR
};

inline std::ostream& operator<<(std::ostream& s,
                                LinearSolverTerminationType type) {
  switch (type) {
    case LinearSolverTerminationType::SUCCESS:
      s << "LINEAR_SOLVER_SUCCESS";
      break;
    case LinearSolverTerminationType::NO_CONVERGENCE:
      s << "LINEAR_SOLVER_NO_CONVERGENCE";
      break;
    case LinearSolverTerminationType::FAILURE:
      s << "LINEAR_SOLVER_FAILURE";
      break;
    case LinearSolverTerminationType::FATAL_ERROR:
      s << "LINEAR_SOLVER_FATAL_ERROR";
      break;
    default:
      s << "UNKNOWN LinearSolverTerminationType";
  }
  return s;
}

// This enum controls the fill-reducing ordering a sparse linear
// algebra library should use before computing a sparse factorization
// (usually Cholesky).
//
// TODO(sameeragarwal): Add support for nested dissection
enum class OrderingType {
  NATURAL,  // Do not re-order the matrix. This is useful when the
            // matrix has been ordered using a fill-reducing ordering
            // already.

  AMD,  // Use the Approximate Minimum Degree algorithm to re-order
        // the matrix.

  NESDIS,  // Use the Nested Dissection algorithm to re-order the matrix.
};

inline std::ostream& operator<<(std::ostream& s, OrderingType type) {
  switch (type) {
    case OrderingType::NATURAL:
      s << "NATURAL";
      break;
    case OrderingType::AMD:
      s << "AMD";
      break;
    case OrderingType::NESDIS:
      s << "NESDIS";
      break;
    default:
      s << "UNKNOWN OrderingType";
  }
  return s;
}

class LinearOperator;

// Abstract base class for objects that implement algorithms for
// solving linear systems
//
//   Ax = b
//
// It is expected that a single instance of a LinearSolver object
// maybe used multiple times for solving multiple linear systems with
// the same sparsity structure. This allows them to cache and reuse
// information across solves. This means that calling Solve on the
// same LinearSolver instance with two different linear systems will
// result in undefined behaviour.
//
// Subclasses of LinearSolver use two structs to configure themselves.
// The Options struct configures the LinearSolver object for its
// lifetime. The PerSolveOptions struct is used to specify options for
// a particular Solve call.
class CERES_NO_EXPORT LinearSolver {
 public:
  struct Options {
    LinearSolverType type = SPARSE_NORMAL_CHOLESKY;
    PreconditionerType preconditioner_type = JACOBI;
    VisibilityClusteringType visibility_clustering_type = CANONICAL_VIEWS;
    DenseLinearAlgebraLibraryType dense_linear_algebra_library_type = EIGEN;
    SparseLinearAlgebraLibraryType sparse_linear_algebra_library_type =
        SUITE_SPARSE;
    OrderingType ordering_type = OrderingType::NATURAL;

    // See solver.h for information about these flags.
    bool dynamic_sparsity = false;
    bool use_explicit_schur_complement = false;

    // Number of internal iterations that the solver uses. This
    // parameter only makes sense for iterative solvers like CG.
    int min_num_iterations = 1;
    int max_num_iterations = 1;

    // Maximum number of iterations performed by SCHUR_POWER_SERIES_EXPANSION.
    // This value controls the maximum number of iterations whether it is used
    // as a preconditioner or just to initialize the solution for
    // ITERATIVE_SCHUR.
    int max_num_spse_iterations = 5;

    // Use SCHUR_POWER_SERIES_EXPANSION to initialize the solution for
    // ITERATIVE_SCHUR. This option can be set true regardless of what
    // preconditioner is being used.
    bool use_spse_initialization = false;

    // When use_spse_initialization is true, this parameter along with
    // max_num_spse_iterations controls the number of
    // SCHUR_POWER_SERIES_EXPANSION iterations performed for initialization. It
    // is not used to control the preconditioner.
    double spse_tolerance = 0.1;

    // If possible, how many threads can the solver use.
    int num_threads = 1;

    // Hints about the order in which the parameter blocks should be
    // eliminated by the linear solver.
    //
    // For example if elimination_groups is a vector of size k, then
    // the linear solver is informed that it should eliminate the
    // parameter blocks 0 ... elimination_groups[0] - 1 first, and
    // then elimination_groups[0] ... elimination_groups[1] - 1 and so
    // on. Within each elimination group, the linear solver is free to
    // choose how the parameter blocks are ordered. Different linear
    // solvers have differing requirements on elimination_groups.
    //
    // The most common use is for Schur type solvers, where there
    // should be at least two elimination groups and the first
    // elimination group must form an independent set in the normal
    // equations. The first elimination group corresponds to the
    // num_eliminate_blocks in the Schur type solvers.
    std::vector<int> elimination_groups;

    // Iterative solvers, e.g. Preconditioned Conjugate Gradients
    // maintain a cheap estimate of the residual which may become
    // inaccurate over time. Thus for non-zero values of this
    // parameter, the solver can be told to recalculate the value of
    // the residual using a |b - Ax| evaluation.
    int residual_reset_period = 10;

    // If the block sizes in a BlockSparseMatrix are fixed, then in
    // some cases the Schur complement based solvers can detect and
    // specialize on them.
    //
    // It is expected that these parameters are set programmatically
    // rather than manually.
    //
    // Please see schur_complement_solver.h and schur_eliminator.h for
    // more details.
    int row_block_size = Eigen::Dynamic;
    int e_block_size = Eigen::Dynamic;
    int f_block_size = Eigen::Dynamic;

    bool use_mixed_precision_solves = false;
    int max_num_refinement_iterations = 0;
    int subset_preconditioner_start_row_block = -1;
    ContextImpl* context = nullptr;
  };

  // Options for the Solve method.
  struct PerSolveOptions {
    // This option only makes sense for unsymmetric linear solvers
    // that can solve rectangular linear systems.
    //
    // Given a matrix A, an optional diagonal matrix D as a vector,
    // and a vector b, the linear solver will solve for
    //
    //   | A | x = | b |
    //   | D |     | 0 |
    //
    // If D is null, then it is treated as zero, and the solver returns
    // the solution to
    //
    //   A x = b
    //
    // In either case, x is the vector that solves the following
    // optimization problem.
    //
    //   arg min_x ||Ax - b||^2 + ||Dx||^2
    //
    // Here A is a matrix of size m x n, with full column rank. If A
    // does not have full column rank, the results returned by the
    // solver cannot be relied on. D, if it is not null is an array of
    // size n.  b is an array of size m and x is an array of size n.
    double* D = nullptr;

    // This option only makes sense for iterative solvers.
    //
    // In general the performance of an iterative linear solver
    // depends on the condition number of the matrix A. For example
    // the convergence rate of the conjugate gradients algorithm
    // is proportional to the square root of the condition number.
    //
    // One particularly useful technique for improving the
    // conditioning of a linear system is to precondition it. In its
    // simplest form a preconditioner is a matrix M such that instead
    // of solving Ax = b, we solve the linear system AM^{-1} y = b
    // instead, where M is such that the condition number k(AM^{-1})
    // is smaller than the conditioner k(A). Given the solution to
    // this system, x = M^{-1} y. The iterative solver takes care of
    // the mechanics of solving the preconditioned system and
    // returning the corrected solution x. The user only needs to
    // supply a linear operator.
    //
    // A null preconditioner is equivalent to an identity matrix being
    // used a preconditioner.
    LinearOperator* preconditioner = nullptr;

    // The following tolerance related options only makes sense for
    // iterative solvers. Direct solvers ignore them.

    // Solver terminates when
    //
    //   |Ax - b| <= r_tolerance * |b|.
    //
    // This is the most commonly used termination criterion for
    // iterative solvers.
    double r_tolerance = 0.0;

    // For PSD matrices A, let
    //
    //   Q(x) = x'Ax - 2b'x
    //
    // be the cost of the quadratic function defined by A and b. Then,
    // the solver terminates at iteration i if
    //
    //   i * (Q(x_i) - Q(x_i-1)) / Q(x_i) < q_tolerance.
    //
    // This termination criterion is more useful when using CG to
    // solve the Newton step. This particular convergence test comes
    // from Stephen Nash's work on truncated Newton
    // methods. References:
    //
    //   1. Stephen G. Nash & Ariela Sofer, Assessing A Search
    //      Direction Within A Truncated Newton Method, Operation
    //      Research Letters 9(1990) 219-221.
    //
    //   2. Stephen G. Nash, A Survey of Truncated Newton Methods,
    //      Journal of Computational and Applied Mathematics,
    //      124(1-2), 45-59, 2000.
    //
    double q_tolerance = 0.0;
  };

  // Summary of a call to the Solve method. We should move away from
  // the true/false method for determining solver success. We should
  // let the summary object do the talking.
  struct Summary {
    double residual_norm = -1.0;
    int num_iterations = -1;
    LinearSolverTerminationType termination_type =
        LinearSolverTerminationType::FAILURE;
    std::string message;
  };

  // If the optimization problem is such that there are no remaining
  // e-blocks, a Schur type linear solver cannot be used. If the
  // linear solver is of Schur type, this function implements a policy
  // to select an alternate nearest linear solver to the one selected
  // by the user. The input linear_solver_type is returned otherwise.
  static LinearSolverType LinearSolverForZeroEBlocks(
      LinearSolverType linear_solver_type);

  virtual ~LinearSolver();

  // Solve Ax = b.
  virtual Summary Solve(LinearOperator* A,
                        const double* b,
                        const PerSolveOptions& per_solve_options,
                        double* x) = 0;

  // This method returns copies instead of references so that the base
  // class implementation does not have to worry about life time
  // issues. Further, this calls are not expected to be frequent or
  // performance sensitive.
  virtual std::map<std::string, CallStatistics> Statistics() const {
    return {};
  }

  // Factory
  static std::unique_ptr<LinearSolver> Create(const Options& options);
};

// This templated subclass of LinearSolver serves as a base class for
// other linear solvers that depend on the particular matrix layout of
// the underlying linear operator. For example some linear solvers
// need low level access to the TripletSparseMatrix implementing the
// LinearOperator interface. This class hides those implementation
// details behind a private virtual method, and has the Solve method
// perform the necessary upcasting.
template <typename MatrixType>
class TypedLinearSolver : public LinearSolver {
 public:
  LinearSolver::Summary Solve(
      LinearOperator* A,
      const double* b,
      const LinearSolver::PerSolveOptions& per_solve_options,
      double* x) override {
    ScopedExecutionTimer total_time("LinearSolver::Solve", &execution_summary_);
    CHECK(A != nullptr);
    CHECK(b != nullptr);
    CHECK(x != nullptr);
    return SolveImpl(down_cast<MatrixType*>(A), b, per_solve_options, x);
  }

  std::map<std::string, CallStatistics> Statistics() const override {
    return execution_summary_.statistics();
  }

 private:
  virtual LinearSolver::Summary SolveImpl(
      MatrixType* A,
      const double* b,
      const LinearSolver::PerSolveOptions& per_solve_options,
      double* x) = 0;

  ExecutionSummary execution_summary_;
};

// Linear solvers that depend on access to the low level structure of
// a SparseMatrix.
// clang-format off
using BlockSparseMatrixSolver = TypedLinearSolver<BlockSparseMatrix>;                  // NOLINT
using CompressedRowSparseMatrixSolver = TypedLinearSolver<CompressedRowSparseMatrix>;  // NOLINT
using DenseSparseMatrixSolver = TypedLinearSolver<DenseSparseMatrix>;                  // NOLINT
using TripletSparseMatrixSolver = TypedLinearSolver<TripletSparseMatrix>;              // NOLINT
// clang-format on

}  // namespace ceres::internal

#include "ceres/internal/reenable_warnings.h"

#endif  // CERES_INTERNAL_LINEAR_SOLVER_H_