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

// spherical_functional.hpp
//
// Energy and gradient of the spherical discrete conformal functional
// evaluated on a ConformalMesh (CGAL::Surface_mesh).
//
// Ported from de.varylab.discreteconformal.functional.SphericalFunctional.
//
// ┌──────────────────────────────────────────────────────────────────────────┐
// │  DOFs                                                                    │
// │    x[v_idx[v]]  = u_v   – conformal factor at vertex v                  │
// │    x[e_idx[e]]  = λ_e   – edge log-length variable (optional)           │
// │    -1 means "pinned" (u_v = 0 / λ_e = λ°_e fixed)                      │
// │                                                                          │
// │  Effective log-length:  Λ_ij = λ°_ij + u_i + u_j                       │
// │  Spherical arc length:  l_ij = 2·asin(min(exp(Λ_ij/2), 1))             │
// │                                                                          │
// │  Gradient:                                                               │
// │    ∂E/∂u_v  = Θ_v  –  Σ_{faces adj. v} α_v(face)                       │
// │    ∂E/∂λ_e  = α_opp(face⁺) + α_opp(face⁻) – π                         │
// │                                                                          │
// │  Energy:                                                                 │
// │    Computed as the Schläfli path integral                                │
// │      E(x) = ∫₀¹ ⟨G(tx), x⟩ dt                                          │
// │    using 10-point Gauss-Legendre quadrature. This is mathematically      │
// │    exact for any conservative (curl-free) gradient G and is numerically  │
// │    accurate to ~10⁻¹⁰ for smooth angle functions. The gradient check     │
// │    therefore tests curl-freeness of G, which is the key integrability    │
// │    condition for the spherical discrete conformal functional.             │
// └──────────────────────────────────────────────────────────────────────────┘

#include "conformal_mesh.hpp"
#include "spherical_geometry.hpp"
#include <CGAL/boost/graph/iterator.h>
#include <vector>
#include <cmath>
#include <cstdint>

namespace conformallab {

// ── Property-map type aliases ─────────────────────────────────────────────────

/// Property map vertex → `double` for the Spherical functional.
using SpherVMapD = ConformalMesh::Property_map<Vertex_index, double>;
/// Property map vertex → `int`    for the Spherical functional.
using SpherVMapI = ConformalMesh::Property_map<Vertex_index, int>;
/// Property map edge   → `double` for the Spherical functional.
using SpherEMapD = ConformalMesh::Property_map<Edge_index,   double>;
/// Property map edge   → `int`    for the Spherical functional.
using SpherEMapI = ConformalMesh::Property_map<Edge_index,   int>;

// ── Persistent map bundle ─────────────────────────────────────────────────────

/// Bundle of the five property maps consumed by the Spherical functional.
struct SphericalMaps {
    SpherVMapI v_idx;    ///< DOF index per vertex (−1 = pinned / u_v = 0).
    SpherEMapI e_idx;    ///< DOF index per edge   (−1 = no edge DOF).
    SpherVMapD theta_v;  ///< Target cone angle Θᵥ (default 2π).
    SpherEMapD theta_e;  ///< Target edge angle θₑ (default π).
    SpherEMapD lambda0;  ///< Base log-length λ⁰ₑ (default 0).
};

// Defaults: theta_v = 2π, theta_e = π, lambda0 = 0.
// lambda0 = 0 means exp(λ°/2)=1, i.e., l=π — degenerate unless u_i<0.
/// Attach the five spherical property maps to `mesh` and return their
/// handles.  Mirrors `setup_euclidean_maps` but uses the `"sv:"` /
/// `"se:"` prefix so the two functionals can coexist on the same mesh
/// (useful for cross-validation tests).
///
/// Defaults match the Euclidean defaults except that `lambda0 = 0` here
/// gives `l_e = π` which is degenerate on the unit sphere — always call
/// `compute_lambda0_from_mesh(mesh, m)` next on a real mesh.
inline SphericalMaps setup_spherical_maps(ConformalMesh& mesh)
{
    SphericalMaps m;
    m.v_idx  = mesh.add_property_map<Vertex_index, int>   ("sv:idx",    -1     ).first;
    m.e_idx  = mesh.add_property_map<Edge_index,   int>   ("se:idx",    -1     ).first;
    m.theta_v= mesh.add_property_map<Vertex_index, double>("sv:theta",  2.0*PI_SPHER).first;
    m.theta_e= mesh.add_property_map<Edge_index,   double>("se:theta",  PI_SPHER    ).first;
    m.lambda0= mesh.add_property_map<Edge_index,   double>("se:lam0",   0.0    ).first;
    return m;
}

/// Assign sequential DOF indices `0..n-1` to all vertices (no edge DOFs).
/// Caller is expected to pin one gauge vertex with `m.v_idx[v] = -1`.
inline int assign_vertex_dof_indices(ConformalMesh& mesh, SphericalMaps& m)
{
    int idx = 0;
    for (auto v : mesh.vertices()) m.v_idx[v] = idx++;
    return idx;
}

/// Assign DOF indices for all vertices AND all edges (vertex-DOFs first,
/// then edge-DOFs).  Mirrors `assign_euclidean_all_dof_indices` for the
/// cyclic spherical formulation.
inline int assign_all_spherical_dof_indices(ConformalMesh& mesh, SphericalMaps& m)
{
    int idx = 0;
    for (auto v : mesh.vertices()) m.v_idx[v] = idx++;
    for (auto e : mesh.edges())    m.e_idx[e] = idx++;
    return idx;
}

/// Count the free DOFs (vertices + edges with index `≥ 0`).
inline int spherical_dimension(const ConformalMesh& mesh, const SphericalMaps& m)
{
    int dim = 0;
    for (auto v : mesh.vertices()) if (m.v_idx[v] >= 0) ++dim;
    for (auto e : mesh.edges())    if (m.e_idx[e] >= 0) ++dim;
    return dim;
}

/// Compute `λ°_e` for every edge from the input vertex positions,
/// assuming `mesh` has vertices on the unit sphere.
///
/// Formula: `λ°_e = 2·log(sin(l_e / 2))` where `l_e = arccos(p_i · p_j)`
/// is the spherical arc length of edge `e`.
///
/// \pre  Every vertex `v` of `mesh` lies on the unit sphere (norm = 1).
/// \pre  No edge is degenerate (`p_i ≠ p_j` and `p_i ≠ -p_j`).
inline void compute_lambda0_from_mesh(ConformalMesh& mesh, SphericalMaps& m)
{
    for (auto e : mesh.edges()) {
        auto h  = mesh.halfedge(e);
        auto p1 = mesh.point(mesh.source(h));
        auto p2 = mesh.point(mesh.target(h));
        // Dot product (works for unit-sphere vertices).
        double dot = p1.x()*p2.x() + p1.y()*p2.y() + p1.z()*p2.z();
        dot = std::max(-1.0, std::min(1.0, dot));
        double l_e = std::acos(dot);              // spherical arc length
        double half_sin = std::sin(l_e * 0.5);   // = exp(λ°/2)
        if (half_sin > 1e-15)
            m.lambda0[e] = 2.0 * std::log(half_sin);
        else
            m.lambda0[e] = -30.0;  // very short edge: essentially 0
    }
}

// ── Evaluation result ─────────────────────────────────────────────────────────

/// Output of `evaluate_spherical()` — energy plus optional gradient.
struct SphericalResult {
    double              energy   = 0.0; ///< Functional value at input DOFs.
    std::vector<double> gradient;       ///< Gradient ∇E (empty if not requested).
};

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

/// Read DOF value from `x` for index `idx`; return 0 if pinned (idx < 0).
static inline double spher_dof_val(int idx, const std::vector<double>& x)
{
    return idx >= 0 ? x[static_cast<std::size_t>(idx)] : 0.0;
}

/// Convert a CGAL half-edge index to a plain `std::size_t` for vector indexing.
static inline std::size_t spher_hidx(Halfedge_index h)
{
    return static_cast<std::size_t>(static_cast<std::uint32_t>(h));
}

// ── Gradient only (no energy) ─────────────────────────────────────────────────

/// Compute the Spherical-functional gradient G(x):
///   * `G_v = Θ_v − Σ_faces α_v(face)`
///   * `G_e = α_opp(face⁺) + α_opp(face⁻) − θ_e`
///
/// The corner angle α_v is stored on half-edges via the convention
/// `h_alpha[h] = corner angle at the vertex ACROSS FROM the edge of h
/// in its face`, which makes both gradient accumulators natural.
inline std::vector<double> spherical_gradient(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const SphericalMaps&       m)
{
    const int n = spherical_dimension(mesh, m);
    std::vector<double> G(static_cast<std::size_t>(n), 0.0);

    // Temporary per-halfedge corner-angle storage.
    // h_alpha[h] = corner angle at the vertex opposite to the edge of h.
    const std::size_t nh = mesh.number_of_halfedges();
    std::vector<double> h_alpha(nh, 0.0);

    // ── Pass 1: compute corner angles per face ────────────────────────────────
    for (auto f : mesh.faces()) {
        Halfedge_index h0 = mesh.halfedge(f);
        Halfedge_index h1 = mesh.next(h0);
        Halfedge_index h2 = mesh.next(h1);

        Vertex_index v1 = mesh.source(h0);
        Vertex_index v2 = mesh.source(h1);
        Vertex_index v3 = mesh.source(h2);

        Edge_index e12 = mesh.edge(h0);
        Edge_index e23 = mesh.edge(h1);
        Edge_index e31 = mesh.edge(h2);

        // Effective log-length Λ_ij = λ°_ij + u_i + u_j
        double u1 = spher_dof_val(m.v_idx[v1], x);
        double u2 = spher_dof_val(m.v_idx[v2], x);
        double u3 = spher_dof_val(m.v_idx[v3], x);

        double lam12 = m.lambda0[e12] + u1 + u2 + spher_dof_val(m.e_idx[e12], x);
        double lam23 = m.lambda0[e23] + u2 + u3 + spher_dof_val(m.e_idx[e23], x);
        double lam31 = m.lambda0[e31] + u3 + u1 + spher_dof_val(m.e_idx[e31], x);

        double l12 = spherical_l(lam12);
        double l23 = spherical_l(lam23);
        double l31 = spherical_l(lam31);

        SphericalFaceAngles fa = spherical_angles(l12, l23, l31);

        if (!fa.valid) continue;   // degenerate face: contributes 0

        // Store convention: h_alpha[h] = corner angle at source(prev(h))
        //   h0 (e12): opposite vertex is v3 → source(prev(h0)) = source(h2) = v3 → α3
        //   h1 (e23): opposite vertex is v1 → source(prev(h1)) = source(h0) = v1 → α1
        //   h2 (e31): opposite vertex is v2 → source(prev(h2)) = source(h1) = v2 → α2
        h_alpha[spher_hidx(h0)] = fa.alpha3;
        h_alpha[spher_hidx(h1)] = fa.alpha1;
        h_alpha[spher_hidx(h2)] = fa.alpha2;
    }

    // ── Pass 2: accumulate gradient ───────────────────────────────────────────

    // Vertex: G_v = Θ_v − Σ h_alpha[prev(h)] for each incoming non-border h to v.
    for (auto v : mesh.vertices()) {
        int iv = m.v_idx[v];
        if (iv < 0) continue;
        double sum_alpha = 0.0;
        for (auto h : CGAL::halfedges_around_target(v, mesh)) {
            if (mesh.is_border(h)) continue;
            sum_alpha += h_alpha[spher_hidx(mesh.prev(h))];
        }
        G[static_cast<std::size_t>(iv)] = m.theta_v[v] - sum_alpha;
    }

    // Edge: G_e for λ_e additive (Λ_ij = λ°_ij + u_i + u_j + λ_e).
    //
    // From the Schläfli identity applied to the spherical face,
    // the contribution of edge DOF λ_e from face f is:
    //   a_f = (2·α_opp − S_f) / 2    where S_f = Σ angles in face f.
    //
    // Summing over both adjacent faces:
    //   G_e = a_f+ + a_f−
    //       = α_opp⁺ + α_opp⁻ − (S_f⁺ + S_f⁻) / 2 − θ_e
    //
    // For flat (Euclidean) triangles S_f = π, recovering the familiar
    // α_opp⁺ + α_opp⁻ − π formula.  For spherical triangles S_f > π.
    for (auto e : mesh.edges()) {
        int ie = m.e_idx[e];
        if (ie < 0) continue;
        auto h  = mesh.halfedge(e);
        auto ho = mesh.opposite(h);
        double sum = 0.0;
        if (!mesh.is_border(h)) {
            double alpha_opp = h_alpha[spher_hidx(h)];
            double S_f = alpha_opp
                       + h_alpha[spher_hidx(mesh.next(h))]
                       + h_alpha[spher_hidx(mesh.prev(h))];
            sum += (2.0 * alpha_opp - S_f) * 0.5;
        }
        if (!mesh.is_border(ho)) {
            double alpha_opp = h_alpha[spher_hidx(ho)];
            double S_f = alpha_opp
                       + h_alpha[spher_hidx(mesh.next(ho))]
                       + h_alpha[spher_hidx(mesh.prev(ho))];
            sum += (2.0 * alpha_opp - S_f) * 0.5;
        }
        G[static_cast<std::size_t>(ie)] = sum - m.theta_e[e];
    }

    return G;
}

// ── Energy via Gauss-Legendre path integral ───────────────────────────────────
//
// E(x) = ∫₀¹ ⟨G(tx), x⟩ dt
//
// This is the correct potential for any conservative G = ∇E.
// Uses 10-point Gauss-Legendre quadrature; error ≈ O(h²⁰) for smooth G.
//
// 10-point GL nodes and weights on [0, 1] (transformed from [-1, 1]):
//   t_k = (1 + s_k) / 2,  w_k = w_GL_k / 2
/// Spherical energy `E(x) = ∫₀¹ ⟨G(t·x), x⟩ dt`, evaluated with
/// 10-point Gauss-Legendre quadrature.  This is the correct potential
/// for any conservative `G = ∇E`; error ≈ O(h²⁰) for smooth G.
inline double spherical_energy(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const SphericalMaps&       m)
{
    // 10-point Gauss-Legendre nodes and weights on [-1, 1].
    static const double gl_s[10] = {
        -0.9739065285171717,  -0.8650633666889845,
        -0.6794095682990244,  -0.4333953941292472,
        -0.1488743389816312,   0.1488743389816312,
         0.4333953941292472,   0.6794095682990244,
         0.8650633666889845,   0.9739065285171717
    };
    static const double gl_w[10] = {
        0.0666713443086881,  0.1494513491505806,
        0.2190863625159820,  0.2692667193099963,
        0.2955242247147529,  0.2955242247147529,
        0.2692667193099963,  0.2190863625159820,
        0.1494513491505806,  0.0666713443086881
    };

    const std::size_t n = x.size();
    double E = 0.0;

    for (int k = 0; k < 10; ++k) {
        double t   = (1.0 + gl_s[k]) * 0.5;   // node on [0, 1]
        double wt  = gl_w[k] * 0.5;            // weight on [0, 1]

        // Evaluate G(t·x)
        std::vector<double> tx(n);
        for (std::size_t i = 0; i < n; ++i) tx[i] = t * x[i];

        auto G = spherical_gradient(mesh, tx, m);

        // Accumulate ⟨G(tx), x⟩ · wt
        double dot = 0.0;
        for (std::size_t i = 0; i < n; ++i) dot += G[i] * x[i];
        E += wt * dot;
    }
    return E;
}

// ── Full evaluation (energy + gradient) ──────────────────────────────────────

/// Evaluate the Spherical functional at DOFs `x`.  Returns energy and
/// gradient (toggle via `need_energy` / `need_gradient`).
inline SphericalResult evaluate_spherical(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const SphericalMaps&       m,
    bool                       need_energy   = true,
    bool                       need_gradient = true)
{
    SphericalResult res;
    if (need_gradient)
        res.gradient = spherical_gradient(mesh, x, m);
    if (need_energy)
        res.energy = spherical_energy(mesh, x, m);
    return res;
}

/// Finite-difference gradient check for the Spherical functional
/// (central differences).  Same defaults as the Java `FunctionalTest`.
inline bool gradient_check_spherical(
    ConformalMesh&             mesh,
    const std::vector<double>& x0,
    const SphericalMaps&       m,
    double                     eps = 1E-5,
    double                     tol = 1E-4)
{
    auto G = spherical_gradient(mesh, x0, m);
    const int n = static_cast<int>(G.size());

    std::vector<double> xp = x0, xm = x0;
    bool ok = true;

    for (int i = 0; i < n; ++i) {
        std::size_t si = static_cast<std::size_t>(i);
        xp[si] = x0[si] + eps;
        xm[si] = x0[si] - eps;

        double Ep = spherical_energy(mesh, xp, m);
        double Em = spherical_energy(mesh, xm, m);

        xp[si] = xm[si] = x0[si];   // restore

        double fd    = (Ep - Em) / (2.0 * eps);
        double err   = std::abs(G[si] - fd);
        double scale = std::max(1.0, std::abs(G[si]));
        if (err / scale > tol) ok = false;
    }
    return ok;
}

// ── Gauge-fix for closed spherical surfaces ───────────────────────────────────
//
// On a closed (boundaryless) spherical surface the energy E(u + t·1) has a
// unique maximum w.r.t. t ∈ ℝ (the "global scale" gauge mode).  Without
// fixing this gauge the functional is unbounded below and no solver converges.
//
// This function returns the scalar shift t* such that
//   Σ_v G_v(x + t*·1_v)  =  0
// i.e., the derivative of E(u+t·1) w.r.t. t is zero at t = t*.
//
// Apply the shift by adding t* to every vertex DOF in x.
//
// Implementation: bisection on f(t) = Σ_v G_v(x + t·1_v).
//   f is strictly monotone decreasing (second derivative < 0) for a convex
//   functional, so bisection converges in O(log₂(2·bracket/tol)) iterations.
//
// Parameters:
//   bracket   – initial search interval [−bracket, +bracket]   (default 50)
//   tol       – absolute tolerance on t*                        (default 1e-8)
//
// Returns 0.0 if the zero cannot be bracketed (already at gauge maximum,
// or open surface — no shift needed).
/// Find the global-scale gauge shift `t*` for the closed-spherical case
/// (see comment block above for the maths).  Apply via `apply_spherical_gauge`.
inline double spherical_gauge_shift(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const SphericalMaps&       m,
    double                     bracket = 50.0,
    double                     tol     = 1e-8)
{
    // Helper: sum of all vertex gradient components at x + t·1_v.
    auto sum_Gv = [&](double t) -> double {
        // Build a shifted copy of x (only vertex DOFs shifted).
        std::vector<double> xt = x;
        for (auto v : mesh.vertices()) {
            int iv = m.v_idx[v];
            if (iv >= 0)
                xt[static_cast<std::size_t>(iv)] += t;
        }
        auto G = spherical_gradient(mesh, xt, m);
        double sum = 0.0;
        for (auto v : mesh.vertices()) {
            int iv = m.v_idx[v];
            if (iv >= 0)
                sum += G[static_cast<std::size_t>(iv)];
        }
        return sum;
    };

    // ── Try bisection first (works when there is a sign change in [−bracket, +bracket]) ──
    double a = -bracket, b = bracket;
    double fa = sum_Gv(a), fb = sum_Gv(b);

    if (fa * fb <= 0.0) {
        for (int iter = 0; iter < 120; ++iter) {
            double c  = 0.5 * (a + b);
            double fc = sum_Gv(c);
            if (std::abs(fc) < tol || (b - a) < tol) return c;
            if (fa * fc < 0.0) { b = c; fb = fc; }
            else               { a = c; fa = fc; }
        }
        return 0.5 * (a + b);
    }

    // ── No sign change (zero may lie at a domain boundary).  ───────────────────
    //    Use damped Newton's method with backtracking line search.
    //    f'(t) estimated by forward finite difference.
    //    When the Newton step overshoots the valid domain (ΣG_v jumps back up
    //    because faces become degenerate), backtracking halves the step until
    //    |f| strictly decreases.
    const double fd_eps = 1e-5;
    double t  = 0.0;
    double ft = sum_Gv(t);

    for (int iter = 0; iter < 120; ++iter) {
        if (std::abs(ft) < tol) return t;

        double ftp = sum_Gv(t + fd_eps);
        double dft = (ftp - ft) / fd_eps;
        if (std::abs(dft) < 1e-14) return t;   // gradient flat — give up

        double dt_raw = -ft / dft;

        // Backtracking line search: halve dt until |f| decreases.
        double alpha = 1.0;
        bool improved = false;
        for (int back = 0; back < 40; ++back) {
            double t_try  = t + alpha * dt_raw;
            t_try = std::max(-bracket, std::min(bracket, t_try));
            double ft_try = sum_Gv(t_try);
            if (std::abs(ft_try) < std::abs(ft)) {
                t  = t_try;
                ft = ft_try;
                improved = true;
                break;
            }
            alpha *= 0.5;
        }
        if (!improved) return t;   // cannot reduce further
    }
    return t;
}

/// Apply the spherical gauge shift in-place: `x_v ← x_v + t*` for every
/// variable vertex, where `t* = spherical_gauge_shift(mesh, x, m, ...)`.
inline void apply_spherical_gauge(
    ConformalMesh&       mesh,
    std::vector<double>& x,
    const SphericalMaps& m,
    double               bracket = 50.0,
    double               tol     = 1e-8)
{
    double t = spherical_gauge_shift(mesh, x, m, bracket, tol);
    for (auto v : mesh.vertices()) {
        int iv = m.v_idx[v];
        if (iv >= 0)
            x[static_cast<std::size_t>(iv)] += t;
    }
}

} // namespace conformallab