// Copyright (c) Sleipnir contributors

#pragma once

#include <chrono>
#include <cmath>
#include <functional>
#include <span>

#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <gch/small_vector.hpp>

#include "sleipnir/optimization/solver/exit_status.hpp"
#include "sleipnir/optimization/solver/iteration_info.hpp"
#include "sleipnir/optimization/solver/newton_matrix_callbacks.hpp"
#include "sleipnir/optimization/solver/options.hpp"
#include "sleipnir/optimization/solver/util/all_finite.hpp"
#include "sleipnir/optimization/solver/util/filter.hpp"
#include "sleipnir/optimization/solver/util/kkt_error.hpp"
#include "sleipnir/optimization/solver/util/regularized_ldlt.hpp"
#include "sleipnir/util/assert.hpp"
#include "sleipnir/util/print_diagnostics.hpp"
#include "sleipnir/util/profiler.hpp"
#include "sleipnir/util/scope_exit.hpp"
#include "sleipnir/util/symbol_exports.hpp"

// See docs/algorithms.md#Works_cited for citation definitions.

namespace slp {

/// Finds the optimal solution to a nonlinear program using Newton's method.
///
/// A nonlinear program has the form:
///
/// ```
/// min_x f(x)
/// ```
///
/// where f(x) is the cost function.
///
/// @tparam Scalar Scalar type.
/// @param[in] matrix_callbacks Matrix callbacks.
/// @param[in] iteration_callbacks The list of callbacks to call at the
///     beginning of each iteration.
/// @param[in] options Solver options.
/// @param[in,out] x The initial guess and output location for the decision
///     variables.
/// @return The exit status.
template <typename Scalar>
ExitStatus newton(
    const NewtonMatrixCallbacks<Scalar>& matrix_callbacks,
    std::span<std::function<bool(const IterationInfo<Scalar>& info)>>
        iteration_callbacks,
    const Options& options, Eigen::Vector<Scalar, Eigen::Dynamic>& x) {
  using DenseVector = Eigen::Vector<Scalar, Eigen::Dynamic>;
  using SparseMatrix = Eigen::SparseMatrix<Scalar>;
  using SparseVector = Eigen::SparseVector<Scalar>;

  using std::isfinite;

  const auto solve_start_time = std::chrono::steady_clock::now();

  gch::small_vector<SolveProfiler> solve_profilers;
  solve_profilers.emplace_back("solver");
  solve_profilers.emplace_back("↳ setup");
  solve_profilers.emplace_back("↳ iteration");
  solve_profilers.emplace_back("  ↳ feasibility check");
  solve_profilers.emplace_back("  ↳ callbacks");
  solve_profilers.emplace_back("  ↳ KKT matrix decomp");
  solve_profilers.emplace_back("  ↳ KKT system solve");
  solve_profilers.emplace_back("  ↳ line search");
  solve_profilers.emplace_back("  ↳ f(x)");
  solve_profilers.emplace_back("  ↳ ∇f(x)");
  solve_profilers.emplace_back("  ↳ ∇²ₓₓL");

  auto& solver_prof = solve_profilers[0];
  auto& setup_prof = solve_profilers[1];
  auto& inner_iter_prof = solve_profilers[2];
  auto& feasibility_check_prof = solve_profilers[3];
  auto& iter_callbacks_prof = solve_profilers[4];
  auto& kkt_matrix_decomp_prof = solve_profilers[5];
  auto& kkt_system_solve_prof = solve_profilers[6];
  auto& line_search_prof = solve_profilers[7];

  // Set up profiled matrix callbacks
#ifndef SLEIPNIR_DISABLE_DIAGNOSTICS
  auto& f_prof = solve_profilers[8];
  auto& g_prof = solve_profilers[9];
  auto& H_prof = solve_profilers[10];

  NewtonMatrixCallbacks<Scalar> matrices{
      matrix_callbacks.num_decision_variables,
      [&](const DenseVector& x) -> Scalar {
        ScopedProfiler prof{f_prof};
        return matrix_callbacks.f(x);
      },
      [&](const DenseVector& x) -> SparseVector {
        ScopedProfiler prof{g_prof};
        return matrix_callbacks.g(x);
      },
      [&](const DenseVector& x) -> SparseMatrix {
        ScopedProfiler prof{H_prof};
        return matrix_callbacks.H(x);
      },
      matrix_callbacks.scaling};
#else
  const auto& matrices = matrix_callbacks;
#endif

  solver_prof.start();
  setup_prof.start();

  Scalar f = matrices.f(x);
  SparseVector g = matrices.g(x);
  SparseMatrix H = matrices.H(x);

  // Ensure matrix callback dimensions are consistent
  slp_assert(g.rows() == matrices.num_decision_variables);
  slp_assert(H.rows() == matrices.num_decision_variables);
  slp_assert(H.cols() == matrices.num_decision_variables);

  DenseVector trial_x;

  Scalar trial_f;

  // Check whether initial guess has finite cost and derivatives
  if (!isfinite(f) || !all_finite(g) || !all_finite(H)) {
    return ExitStatus::NONFINITE_INITIAL_GUESS;
  }

  int iterations = 0;

  Filter<Scalar> filter;

  RegularizedLDLT<Scalar> solver{
      // Use sparse solver if lower triangle fills < 25% of system
      H.nonZeros() < 0.25 * H.size(), matrices.num_decision_variables, 0};

  // Variables for determining when a step is acceptable
  constexpr Scalar α_reduction_factor(0.5);
  constexpr Scalar α_min(1e-20);

  // Error
  Scalar E_0 = unscaled_kkt_error<Scalar, KKTErrorType::INF_NORM_SCALED>(
      matrices.scaling, g);

  setup_prof.stop();

  // Prints final solver diagnostics when the solver exits
  scope_exit exit{[&] {
    if (options.diagnostics) {
      solver_prof.stop();
      if (iterations > 0) {
        print_bottom_iteration_diagnostics();
      }
      print_solver_diagnostics(solve_profilers);
    }
  }};

  while (E_0 > Scalar(options.tolerance)) {
    ScopedProfiler inner_iter_profiler{inner_iter_prof};
    ScopedProfiler feasibility_check_profiler{feasibility_check_prof};

    // Check for diverging iterates
    if (x.template lpNorm<Eigen::Infinity>() > Scalar(1e10) || !x.allFinite()) {
      return ExitStatus::DIVERGING_ITERATES;
    }

    feasibility_check_profiler.stop();
    ScopedProfiler iter_callbacks_profiler{iter_callbacks_prof};

    // Call iteration callbacks
    for (const auto& callback : iteration_callbacks) {
      if (callback({iterations, x, {}, {}, {}, g, H, {}, {}})) {
        return ExitStatus::CALLBACK_REQUESTED_STOP;
      }
    }

    iter_callbacks_profiler.stop();
    ScopedProfiler kkt_matrix_decomp_profiler{kkt_matrix_decomp_prof};

    // Solve the Newton-KKT system
    //
    // Hpˣ = −∇f
    solver.compute(H);

    kkt_matrix_decomp_profiler.stop();
    ScopedProfiler kkt_system_solve_profiler{kkt_system_solve_prof};

    DenseVector p_x = solver.solve(-g);

    kkt_system_solve_profiler.stop();
    ScopedProfiler line_search_profiler{line_search_prof};

    constexpr Scalar α_max(1);
    Scalar α = α_max;

    // Loop until a step is accepted. If a step becomes acceptable, the loop
    // will exit early.
    while (1) {
      trial_x = x + α * p_x;

      trial_f = matrices.f(trial_x);

      // If f(xₖ + αpₖˣ) isn't finite, reduce step size immediately
      if (!isfinite(trial_f)) {
        // Reduce step size
        α *= α_reduction_factor;

        if (α < α_min) {
          return ExitStatus::LINE_SEARCH_FAILED;
        }
        continue;
      }

      // Check whether filter accepts trial iterate
      if (filter.try_add(FilterEntry{f}, FilterEntry{trial_f}, p_x, g, α)) {
        // Accept step
        break;
      }

      // Reduce step size
      α *= α_reduction_factor;

      // If step size hit a minimum, check if the KKT error was reduced. If it
      // wasn't, report bad line search.
      if (α < α_min) {
        Scalar current_kkt_error = kkt_error<Scalar, KKTErrorType::ONE_NORM>(g);

        trial_x = x + α_max * p_x;

        Scalar next_kkt_error =
            kkt_error<Scalar, KKTErrorType::ONE_NORM>(matrices.g(trial_x));

        // If the step using αᵐᵃˣ reduced the KKT error, accept it anyway
        if (next_kkt_error <= Scalar(0.999) * current_kkt_error) {
          trial_f = matrices.f(trial_x);

          // Accept step
          break;
        }

        return ExitStatus::LINE_SEARCH_FAILED;
      }
    }

    line_search_profiler.stop();

    // Update iterates
    x = trial_x;

    f = trial_f;

    // Update autodiff for Hessian
    g = matrices.g(x);
    H = matrices.H(x);

    // Update the error
    E_0 = unscaled_kkt_error<Scalar, KKTErrorType::INF_NORM_SCALED>(
        matrices.scaling, g);

    inner_iter_profiler.stop();

    if (options.diagnostics) {
      print_iteration_diagnostics(
          iterations, IterationType::NORMAL,
          inner_iter_profiler.current_duration(), E_0, f, Scalar(0), Scalar(0),
          Scalar(0), solver.hessian_regularization(),
          solver.constraint_jacobian_regularization(),
          p_x.template lpNorm<Eigen::Infinity>(), Scalar(1), α, α_max,
          α_reduction_factor, Scalar(1));
    }

    ++iterations;

    // Check for max iterations
    if (iterations >= options.max_iterations) {
      return ExitStatus::MAX_ITERATIONS_EXCEEDED;
    }

    // Check for max wall clock time
    if (std::chrono::steady_clock::now() - solve_start_time > options.timeout) {
      return ExitStatus::TIMEOUT;
    }
  }

  return ExitStatus::SUCCESS;
}

extern template SLEIPNIR_DLLEXPORT ExitStatus
newton(const NewtonMatrixCallbacks<double>& matrix_callbacks,
       std::span<std::function<bool(const IterationInfo<double>& info)>>
           iteration_callbacks,
       const Options& options, Eigen::Vector<double, Eigen::Dynamic>& x);

}  // namespace slp
