#pragma once
// Copyright (c) 2024-2026 Tarik Moussa.
// SPDX-License-Identifier: MIT

// newton_solver.hpp
//
// Phase 4a — Newton solver for all three discrete conformal functionals.
//
// Solves G(x) = 0 where G is the gradient of the discrete conformal energy.
//
// ┌──────────────────────────────────────────────────────────────────────────┐
// │  Gradient sign conventions                                               │
// │    Euclidean / Spherical:  G_v = Θ_v − Σ α_v    (target − actual)      │
// │    HyperIdeal:             G_v = Σ β_v − Θ_v    (actual − target)      │
// │                                                                          │
// │  All solvers use the same Newton step  Δx = −H⁻¹·G                     │
// │                                                                          │
// │  Hessian sign at equilibrium                                             │
// │    Euclidean:   H PSD  → SimplicialLDLT on  H                          │
// │    Spherical:   H NSD  → SimplicialLDLT on −H  (solve (−H)Δx = G)      │
// │    HyperIdeal:  H PSD  → SimplicialLDLT on  H  (analytical H: future)  │
// └──────────────────────────────────────────────────────────────────────────┘
//
// SparseQR fallback:
//   When SimplicialLDLT reports a failure (e.g. singular H on a closed mesh
//   without a pinned vertex), the solver automatically retries with
//   Eigen::SparseQR, which finds the minimum-norm Newton step orthogonal to
//   the null space.  This handles the gauge mode on closed surfaces without
//   requiring the caller to pin a vertex explicitly.
//
// Requires:
//   Eigen::SimplicialLDLT, Eigen::SparseQR (Eigen sparse module)

#include "euclidean_hessian.hpp"
#include "spherical_hessian.hpp"
#include "hyper_ideal_hessian.hpp"
#include "cp_euclidean_functional.hpp"
#include "inversive_distance_functional.hpp"
#include <Eigen/SparseCholesky>
#include <Eigen/SparseQR>
#include <Eigen/OrderingMethods>
#include <Eigen/Dense>
#include <algorithm>
#include <cmath>

namespace conformallab {

// ── Result ────────────────────────────────────────────────────────────────────

struct NewtonResult {
    std::vector<double> x;            ///< DOF vector at termination
    int                 iterations;   ///< Newton steps taken
    double              grad_inf_norm;///< max |G_i| at termination
    bool                converged;    ///< true iff grad_inf_norm < tol
};

// ── Internal helpers ──────────────────────────────────────────────────────────

namespace detail {

// Solve A·Δx = rhs with SimplicialLDLT; on failure fall back to SparseQR.
// Returns Δx.  ok is set to false only if both solvers fail.
// If fallback_used is non-null, it is set to true iff SparseQR was needed.
inline Eigen::VectorXd solve_with_fallback(
    const Eigen::SparseMatrix<double>& A,
    const Eigen::VectorXd&             rhs,
    bool&                              ok,
    bool*                              fallback_used = nullptr)
{
    if (fallback_used) *fallback_used = false;

    Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> ldlt(A);
    if (ldlt.info() == Eigen::Success) {
        Eigen::VectorXd dx = ldlt.solve(rhs);
        if (ldlt.info() == Eigen::Success) { ok = true; return dx; }
    }
    // Fallback: SparseQR — handles singular/rank-deficient H (gauge modes).
    if (fallback_used) *fallback_used = true;
    Eigen::SparseQR<Eigen::SparseMatrix<double>, Eigen::COLAMDOrdering<int>> qr(A);
    if (qr.info() == Eigen::Success) {
        Eigen::VectorXd dx = qr.solve(rhs);
        if (qr.info() == Eigen::Success) { ok = true; return dx; }
    }
    ok = false;
    return Eigen::VectorXd::Zero(rhs.size());
}

} // namespace detail

// ── Public linear-system solver (SparseQR fallback) ──────────────────────────
//
// Solve A·x = rhs with Eigen::SimplicialLDLT; if that fails (singular or
// rank-deficient A), retry with Eigen::SparseQR which finds the minimum-norm
// solution orthogonal to the null space.
//
// This is the same primitive used internally by all three Newton solvers.
// Exposing it publicly lets callers (tests, downstream code) reuse the logic
// and — via the optional fallback_used pointer — verify which code path ran.
//
//   fallback_used  – if non-null, set to true iff SparseQR was invoked
//   Returns Eigen::VectorXd::Zero(rhs.size()) if both solvers fail.
inline Eigen::VectorXd solve_linear_system(
    const Eigen::SparseMatrix<double>& A,
    const Eigen::VectorXd&             rhs,
    bool*                              fallback_used = nullptr)
{
    bool ok = false;
    return detail::solve_with_fallback(A, rhs, ok, fallback_used);
}

namespace detail {   // re-open for the remaining helpers

// Backtracking line search: find the largest α in {1, 0.5, 0.25, …} such that
// ||G(x + α·Δx)||₂ < ||G(x)||₂.  Returns the accepted step (α may stay 1).
template <typename GradFn>
inline std::vector<double> line_search(
    const std::vector<double>& x,
    const Eigen::VectorXd&     dx,
    double                     norm0,
    GradFn&&                   grad_fn,
    int                        max_halvings = 20)
{
    const int n = static_cast<int>(x.size());
    double alpha = 1.0;
    std::vector<double> xnew(static_cast<std::size_t>(n));

    for (int ls = 0; ls < max_halvings; ++ls) {
        for (int i = 0; i < n; ++i)
            xnew[static_cast<std::size_t>(i)] = x[static_cast<std::size_t>(i)]
                                                + alpha * dx[i];
        auto Gnew = grad_fn(xnew);
        double norm_new = 0.0;
        for (double v : Gnew) norm_new += v * v;
        norm_new = std::sqrt(norm_new);
        if (norm_new < norm0) return xnew;
        alpha *= 0.5;
    }
    // No improvement found — return best attempt (full step)
    for (int i = 0; i < n; ++i)
        xnew[static_cast<std::size_t>(i)] = x[static_cast<std::size_t>(i)] + dx[i];
    return xnew;
}

} // namespace detail

// ── Euclidean Newton solver ────────────────────────────────────────────────────

/// Solve the Euclidean discrete conformal problem: find u ∈ ℝ^V such that
/// Σ_{faces adj v} α_v(u) = Θ_v for all vertices v.
///
/// Starting from x0, Newton's method minimises E(u) (the Euclidean DCE energy,
/// which is convex) by iterating  u ← u − H⁻¹·G  with backtracking line search.
/// The Hessian H is the cotangent Laplacian — PSD with one zero eigenvalue on
/// closed surfaces (gauge mode).  A SparseQR fallback handles this automatically.
///
/// \param mesh      Input triangulated surface (edges must carry lambda0 + theta_v).
/// \param x0        Initial DOF vector (length = number of free vertices).
///                  Pass all-zeros for a flat start (typical).
/// \param m         EuclideanMaps: lambda0[e], theta_v[v], v_idx[v] must be set.
///                  Call setup_euclidean_maps() + compute_euclidean_lambda0_from_mesh()
///                  + enforce_gauss_bonnet() before passing here.
/// \param tol       Convergence threshold on max |G_i|.  Default: 1e-8.
/// \param max_iter  Maximum Newton iterations.  Default: 200.
/// \return          NewtonResult{x*, iterations, grad_inf_norm, converged}.
///
/// \note On closed meshes without a pinned vertex, SimplicialLDLT detects the
///       gauge singularity and falls back to SparseQR automatically.
///
/// \see doc/math/discrete-conformal-theory.md §3 for the mathematical background.
inline NewtonResult newton_euclidean(
    ConformalMesh&             mesh,
    std::vector<double>        x0,
    const EuclideanMaps&       m,
    double                     tol      = 1e-8,
    int                        max_iter = 200)
{
    std::vector<double> x = x0;
    const int n = static_cast<int>(x.size());

    NewtonResult res;
    res.converged      = false;
    res.iterations     = 0;
    res.grad_inf_norm  = 0.0;

    Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver;

    for (int iter = 0; iter < max_iter; ++iter) {
        // ── Gradient ──────────────────────────────────────────────────────────
        auto G_std = euclidean_gradient(mesh, x, m);
        Eigen::Map<const Eigen::VectorXd> G(G_std.data(), n);

        double inf_norm = G.cwiseAbs().maxCoeff();
        if (inf_norm < tol) {
            res.converged     = true;
            res.grad_inf_norm = inf_norm;
            res.iterations    = iter;
            res.x             = x;
            return res;
        }

        // ── Hessian + solve H·Δx = −G  (SparseQR fallback for singular H) ──
        auto H = euclidean_hessian(mesh, x, m);
        bool ok = false;
        Eigen::VectorXd dx = detail::solve_with_fallback(H, -G, ok);
        if (!ok) break;

        // ── Backtracking line search ──────────────────────────────────────────
        double norm0 = G.norm();
        x = detail::line_search(x, dx, norm0,
            [&](const std::vector<double>& xnew) {
                return euclidean_gradient(mesh, xnew, m);
            });

        res.iterations = iter + 1;
    }

    // Report final gradient norm
    auto G_final = euclidean_gradient(mesh, x, m);
    double inf_final = 0.0;
    for (double v : G_final) inf_final = std::max(inf_final, std::abs(v));
    res.grad_inf_norm = inf_final;
    res.x             = x;
    return res;
}

// ── Spherical Newton solver ───────────────────────────────────────────────────

/// Solve the spherical discrete conformal problem: find u ∈ ℝ^V such that
/// Σ_{faces adj v} α_v(u) = Θ_v for all vertices v (genus-0 / sphere-like surfaces).
///
/// The spherical DCE energy is *concave*, so the Hessian H is NSD at the solution.
/// The solver factorises −H (which is PSD) and solves (−H)·Δx = G.
/// A gauge vertex must be pinned (set v_idx = -1) to remove the rotational mode.
///
/// \param mesh      Input triangulated surface, genus 0.
/// \param x0        Initial DOF vector (length = free vertices, excluding gauge_vertex).
///                  All-zeros is a good start.
/// \param m         SphericalMaps: lambda0[e], theta_v[v], v_idx[v], gauge_vertex set.
///                  Call setup_spherical_maps() + compute_spherical_lambda0_from_mesh()
///                  + enforce_gauss_bonnet() (checks Σ(2π-Θ) > 0) before passing here.
/// \param tol       Convergence threshold on max |G_i|.  Default: 1e-8.
/// \param max_iter  Maximum Newton iterations.  Default: 200.
/// \return          NewtonResult{x*, iterations, grad_inf_norm, converged}.
///
/// \note Unlike the Euclidean solver, the spherical solver does NOT need a SparseQR
///       fallback — the gauge vertex pins the null mode directly.
///
/// \see doc/math/geometry-modes.md §Spherical for sign-convention details.
inline NewtonResult newton_spherical(
    ConformalMesh&             mesh,
    std::vector<double>        x0,
    const SphericalMaps&       m,
    double                     tol      = 1e-8,
    int                        max_iter = 200)
{
    std::vector<double> x = x0;
    const int n = static_cast<int>(x.size());

    NewtonResult res;
    res.converged      = false;
    res.iterations     = 0;
    res.grad_inf_norm  = 0.0;

    for (int iter = 0; iter < max_iter; ++iter) {
        // ── Gradient ──────────────────────────────────────────────────────────
        auto G_std = spherical_gradient(mesh, x, m);
        Eigen::Map<const Eigen::VectorXd> G(G_std.data(), n);

        double inf_norm = G.cwiseAbs().maxCoeff();
        if (inf_norm < tol) {
            res.converged     = true;
            res.grad_inf_norm = inf_norm;
            res.iterations    = iter;
            res.x             = x;
            return res;
        }

        // ── Hessian: negate to get PSD; solve (−H)·Δx = G  ──────────────────
        auto H    = spherical_hessian(mesh, x, m);
        auto negH = Eigen::SparseMatrix<double>(-H);
        bool ok   = false;
        Eigen::VectorXd dx = detail::solve_with_fallback(negH, G, ok);
        if (!ok) break;

        // ── Backtracking line search ──────────────────────────────────────────
        double norm0 = G.norm();
        x = detail::line_search(x, dx, norm0,
            [&](const std::vector<double>& xnew) {
                return spherical_gradient(mesh, xnew, m);
            });

        res.iterations = iter + 1;
    }

    auto G_final = spherical_gradient(mesh, x, m);
    double inf_final = 0.0;
    for (double v : G_final) inf_final = std::max(inf_final, std::abs(v));
    res.grad_inf_norm = inf_final;
    res.x             = x;
    return res;
}

// ── HyperIdeal Newton solver ──────────────────────────────────────────────────

/// Solve the hyper-ideal discrete conformal problem: find (b, a) ∈ ℝ^{V+E} such that
/// Σ β_v(b,a) = Θ_v  and  Σ α_e(b,a) = θ_e  for all vertices v and edges e.
///
/// Used for genus-g surfaces (g ≥ 1) under hyperbolic cone metrics.
/// The energy is *strictly convex* (Springborn 2020, Theorem 1.3), so Newton
/// converges globally from any starting point.
///
/// DOF layout: first V_free entries are vertex variables b_v (hyper-ideal radii),
/// followed by E entries for edge variables a_e (intersection angles).
/// Use assign_all_dof_indices(mesh, maps) to set v_idx and e_idx automatically —
/// no vertex needs to be pinned.
///
/// \param mesh      Input triangulated surface, genus g ≥ 1.
/// \param x0        Initial DOF vector (length = V + E).  All-zeros typical.
/// \param m         HyperIdealMaps: lambda0[e], theta_v[v], v_idx[v], e_idx[e] set.
///                  Call setup_hyper_ideal_maps() + compute_hyper_ideal_lambda0_from_mesh().
/// \param tol       Convergence threshold on max |G_i|.  Default: 1e-8.
/// \param max_iter  Maximum Newton iterations.  Default: 200.
/// \param hess_eps  Finite-difference step for Hessian approximation.  Default: 1e-5.
///                  (Phase 9b will replace this with an analytic Hessian.)
/// \return          NewtonResult{x*, iterations, grad_inf_norm, converged}.
///
/// \see Springborn (2020), Theorem 1.3 for the strict convexity proof.
/// \see doc/math/geometry-modes.md §Hyper-ideal for DOF layout details.
inline NewtonResult newton_hyper_ideal(
    ConformalMesh&             mesh,
    std::vector<double>        x0,
    const HyperIdealMaps&      m,
    double                     tol      = 1e-8,
    int                        max_iter = 200,
    double                     hess_eps = 1e-5)
{
    std::vector<double> x = x0;
    const int n = static_cast<int>(x.size());

    NewtonResult res;
    res.converged      = false;
    res.iterations     = 0;
    res.grad_inf_norm  = 0.0;

    for (int iter = 0; iter < max_iter; ++iter) {
        // ── Gradient ──────────────────────────────────────────────────────────
        auto G_std = evaluate_hyper_ideal(mesh, x, m, /*energy=*/false).gradient;
        Eigen::Map<const Eigen::VectorXd> G(G_std.data(), n);

        double inf_norm = G.cwiseAbs().maxCoeff();
        if (inf_norm < tol) {
            res.converged     = true;
            res.grad_inf_norm = inf_norm;
            res.iterations    = iter;
            res.x             = x;
            return res;
        }

        // ── Hessian (numerical FD) + solve H·Δx = −G ─────────────────────────
        auto H  = hyper_ideal_hessian_sym(mesh, x, m, hess_eps);
        bool ok = false;
        Eigen::VectorXd dx = detail::solve_with_fallback(H, -G, ok);
        if (!ok) break;

        // ── Backtracking line search ──────────────────────────────────────────
        double norm0 = G.norm();
        x = detail::line_search(x, dx, norm0,
            [&](const std::vector<double>& xnew) {
                return evaluate_hyper_ideal(mesh, xnew, m, false).gradient;
            });

        res.iterations = iter + 1;
    }

    auto G_final = evaluate_hyper_ideal(mesh, x, m, false).gradient;
    double inf_final = 0.0;
    for (double v : G_final) inf_final = std::max(inf_final, std::abs(v));
    res.grad_inf_norm = inf_final;
    res.x             = x;
    return res;
}

// ── CP-Euclidean Newton solver (Phase 9a.1) ───────────────────────────────────

/// Solve the CP-Euclidean circle-packing problem: find ρ ∈ ℝ^F such that the
/// per-face angle sums match φ_f at every free face.
///
/// The CP-Euclidean energy (Bobenko-Pinkall-Springborn 2010 §6) is strictly
/// convex on its open domain of validity, so the Hessian H is PSD and the
/// solution is unique up to the gauge mode pinned by `f_idx == −1`.
/// `cp_euclidean_hessian` provides the analytic 2×2-per-edge formula
/// `h_jk = sin θ / (cosh Δρ − cos θ)`; no FD machinery is required.
///
/// \param mesh      Input triangle mesh (closed or with boundary).
/// \param x0        Initial DOF vector (length = number of free faces).
///                  All-zeros is a valid start.
/// \param m         CPEuclideanMaps: f_idx must have one pinned face
///                  (`f_idx[f0] == −1`); theta_e and phi_f set by the caller.
/// \param tol       Convergence threshold on `‖G‖∞`.  Default: 1e-8.
/// \param max_iter  Newton iteration limit.  Default: 200.
/// \return          NewtonResult{x*, iterations, grad_inf_norm, converged}.
///
/// \note  Unlike the Euclidean solver, the CP-Euclidean Hessian is exact
///        (analytic), so the SparseQR fallback only triggers in genuine
///        gauge-singular situations (no pinned face).
/// \see   doc/architecture/phase-9a-validation.md §1 for the BPS-2010 mapping.
inline NewtonResult newton_cp_euclidean(
    ConformalMesh&             mesh,
    std::vector<double>        x0,
    const CPEuclideanMaps&     m,
    double                     tol      = 1e-8,
    int                        max_iter = 200)
{
    std::vector<double> x = x0;
    const int n = static_cast<int>(x.size());

    NewtonResult res;
    res.converged      = false;
    res.iterations     = 0;
    res.grad_inf_norm  = 0.0;

    for (int iter = 0; iter < max_iter; ++iter) {
        auto G_std = cp_euclidean_gradient(mesh, x, m);
        Eigen::Map<const Eigen::VectorXd> G(G_std.data(), n);

        double inf_norm = G.cwiseAbs().maxCoeff();
        if (inf_norm < tol) {
            res.converged     = true;
            res.grad_inf_norm = inf_norm;
            res.iterations    = iter;
            res.x             = x;
            return res;
        }

        auto H = cp_euclidean_hessian(mesh, x, m);
        bool ok = false;
        Eigen::VectorXd dx = detail::solve_with_fallback(H, -G, ok);
        if (!ok) break;

        double norm0 = G.norm();
        x = detail::line_search(x, dx, norm0,
            [&](const std::vector<double>& xnew) {
                return cp_euclidean_gradient(mesh, xnew, m);
            });

        res.iterations = iter + 1;
    }

    auto G_final = cp_euclidean_gradient(mesh, x, m);
    double inf_final = 0.0;
    for (double v : G_final) inf_final = std::max(inf_final, std::abs(v));
    res.grad_inf_norm = inf_final;
    res.x             = x;
    return res;
}

// ── Inversive-Distance Newton solver (Phase 9a.2) ─────────────────────────────

/// Solve the inversive-distance circle-packing problem: find u ∈ ℝ^V such that
/// Σ_{faces adj v} α_v(u) = Θ_v at every free vertex (Luo 2004 Lemma 3.1).
///
/// The inversive-distance energy is (locally) strictly convex on the open
/// domain where every triangle satisfies the inequalities.  Luo's 1-form is
/// closed there, so the path-integral energy is well-defined.
///
/// MVP implementation: the Hessian is computed by **finite differences** of
/// the analytic gradient (same pattern as the Phase 4a HyperIdeal solver).
/// An analytic Hessian via Glickenstein 2011 eq. (4.6) is tracked in
/// `doc/roadmap/research-track.md` as Phase 9a.2-analytic.
///
/// \param mesh      Input triangle mesh.
/// \param x0        Initial DOF vector (length = number of free vertices).
/// \param m         InversiveDistanceMaps: v_idx has at least one pinned
///                  vertex; I_e and r0 set by compute_inversive_distance_init.
/// \param tol       Convergence threshold on `‖G‖∞`.  Default: 1e-8.
/// \param max_iter  Newton iteration limit.  Default: 200.
/// \param hess_eps  FD step size for the Hessian.  Default: 1e-5.
/// \return          NewtonResult{x*, iterations, grad_inf_norm, converged}.
///
/// \note  Convergence is sensitive to the initial point: u = 0 is the
///        natural choice when `compute_inversive_distance_init_from_mesh`
///        has been called, since the Bowers-Stephenson identity reconstructs
///        the input edge lengths at u = 0.
inline NewtonResult newton_inversive_distance(
    ConformalMesh&                mesh,
    std::vector<double>           x0,
    const InversiveDistanceMaps&  m,
    double                        tol      = 1e-8,
    int                           max_iter = 200,
    double                        hess_eps = 1e-5)
{
    std::vector<double> x = x0;
    const int n = static_cast<int>(x.size());

    NewtonResult res;
    res.converged      = false;
    res.iterations     = 0;
    res.grad_inf_norm  = 0.0;

    // Local FD Hessian builder — n × (cost of gradient eval).
    auto build_hessian = [&](const std::vector<double>& xc) -> Eigen::SparseMatrix<double> {
        std::vector<Eigen::Triplet<double>> trips;
        trips.reserve(static_cast<std::size_t>(n) * 16);   // sparse heuristic

        std::vector<double> xp = xc, xm = xc;
        for (int j = 0; j < n; ++j) {
            const std::size_t sj = static_cast<std::size_t>(j);
            xp[sj] = xc[sj] + hess_eps;
            xm[sj] = xc[sj] - hess_eps;

            auto Gp = inversive_distance_gradient(mesh, xp, m);
            auto Gm = inversive_distance_gradient(mesh, xm, m);

            xp[sj] = xm[sj] = xc[sj];   // restore

            for (int i = 0; i < n; ++i) {
                double val = (Gp[static_cast<std::size_t>(i)]
                            - Gm[static_cast<std::size_t>(i)])
                             / (2.0 * hess_eps);
                if (std::abs(val) > 1e-15)
                    trips.emplace_back(i, j, val);
            }
        }
        Eigen::SparseMatrix<double> H(n, n);
        H.setFromTriplets(trips.begin(), trips.end());
        // Symmetrise — FD rounding may introduce tiny asymmetries.
        Eigen::SparseMatrix<double> Ht = H.transpose();
        return (H + Ht) * 0.5;
    };

    for (int iter = 0; iter < max_iter; ++iter) {
        auto G_std = inversive_distance_gradient(mesh, x, m);
        Eigen::Map<const Eigen::VectorXd> G(G_std.data(), n);

        double inf_norm = G.cwiseAbs().maxCoeff();
        if (inf_norm < tol) {
            res.converged     = true;
            res.grad_inf_norm = inf_norm;
            res.iterations    = iter;
            res.x             = x;
            return res;
        }

        auto H = build_hessian(x);
        bool ok = false;
        Eigen::VectorXd dx = detail::solve_with_fallback(H, -G, ok);
        if (!ok) break;

        double norm0 = G.norm();
        x = detail::line_search(x, dx, norm0,
            [&](const std::vector<double>& xnew) {
                return inversive_distance_gradient(mesh, xnew, m);
            });

        res.iterations = iter + 1;
    }

    auto G_final = inversive_distance_gradient(mesh, x, m);
    double inf_final = 0.0;
    for (double v : G_final) inf_final = std::max(inf_final, std::abs(v));
    res.grad_inf_norm = inf_final;
    res.x             = x;
    return res;
}

} // namespace conformallab