#pragma once
// 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 <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 ────────────────────────────────────────────────────
//
// Minimises the Euclidean discrete conformal energy by solving G(x) = 0.
// The Hessian H is PSD; Eigen::SimplicialLDLT is used directly.
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 ───────────────────────────────────────────────────
//
// Solves G(x) = 0 for the spherical discrete conformal functional.
// The Hessian H is NSD at the solution; −H is PSD, so we factorise −H and
// solve (−H)·Δx = G  ⟺  H·Δx = −G.
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 ──────────────────────────────────────────────────
//
// Solves G(x) = 0 for the hyper-ideal discrete conformal functional.
//
// Gradient sign convention (opposite to Euclidean/Spherical):
//   G_v = Σ β_v − Θ_v,   G_e = Σ α_e − θ_e   (actual − target)
//
// The hyper-ideal energy is strictly convex (Springborn 2020), so H is PSD
// and SimplicialLDLT (with SparseQR fallback) applies directly.
//
// The Hessian is computed by symmetric finite differences of G (see
// hyper_ideal_hessian.hpp); replace with an analytical Hessian in Phase 5.
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;
}

} // namespace conformallab