#pragma once
// newton_solver.hpp
//
// Phase 4a — Newton solver for the discrete conformal functionals.
//
// Solves G(x) = 0 where G is the gradient of the discrete conformal energy:
//   G_v = Θ_v − Σ_f α_v^f   (angle-sum residual at each vertex DOF)
//
// Algorithm per iteration:
//   1. Compute gradient G(x)
//   2. Check convergence: max|G_i| < tol → done
//   3. Compute sparse Hessian H(x)
//   4. Factorize and solve the Newton system:
//        Euclidean:  H · Δx = −G     (H is PSD → SimplicialLDLT directly)
//        Spherical:  (−H) · Δx = G   (H is NSD → negate to get PSD matrix)
//   5. Backtracking line search: halve α until ||G(x+α·Δx)|| < ||G(x)||
//   6. x ← x + α·Δx, go to 1
//
// Requires:
//   Eigen::SimplicialLDLT (part of Eigen's sparse Cholesky module)

#include "euclidean_hessian.hpp"
#include "spherical_hessian.hpp"
#include <Eigen/SparseCholesky>
#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 {

// 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 + factorisation ───────────────────────────────────────────
        auto H = euclidean_hessian(mesh, x, m);
        solver.compute(H);
        if (solver.info() != Eigen::Success) break;

        // ── Newton step: solve H·Δx = −G ────────────────────────────────────
        Eigen::VectorXd dx = solver.solve(-G);
        if (solver.info() != Eigen::Success) 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;

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

    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 matrix ────────────────────────────────
        auto H     = spherical_hessian(mesh, x, m);
        auto negH  = Eigen::SparseMatrix<double>(-H);
        solver.compute(negH);
        if (solver.info() != Eigen::Success) break;

        // ── Newton step: solve (−H)·Δx = G  ─────────────────────────────────
        Eigen::VectorXd dx = solver.solve(G);
        if (solver.info() != Eigen::Success) 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;
}

} // namespace conformallab