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

// euclidean_functional.hpp
//
// Energy and gradient of the Euclidean discrete conformal functional
// (EuclideanCyclicFunctional) evaluated on a ConformalMesh.
//
// Ported from de.varylab.discreteconformal.functional.EuclideanCyclicFunctional.
//
// ┌──────────────────────────────────────────────────────────────────────────┐
// │  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 = 0, only λ° contributes)          │
// │                                                                          │
// │  Effective log-length (always additive, unlike SphericalFunctional):    │
// │    Λ̃_ij = λ°_ij + u_i + u_j + (x[e_idx[e]] if variable, else 0)       │
// │                                                                          │
// │  Side length:  l_ij = exp(Λ̃_ij / 2)                                    │
// │                                                                          │
// │  Gradient:                                                               │
// │    ∂E/∂u_v  = Θ_v  −  Σ_{faces adj. v} α_v(face)                       │
// │    ∂E/∂λ_e  = α_opp(face⁺) + α_opp(face⁻) − φ_e                       │
// │                                                                          │
// │  Energy:                                                                 │
// │    Computed as the path integral E(x) = ∫₀¹ ⟨G(tx), x⟩ dt             │
// │    using 10-point Gauss-Legendre quadrature (same as SphericalFunctional)│
// │    This is E(0)=0 by construction and exact for conservative G.         │
// └──────────────────────────────────────────────────────────────────────────┘
//
// Halfedge convention (identical to SphericalFunctional):
//   h0 = mesh.halfedge(f),  h1 = next(h0),  h2 = next(h1)
//   v1 = source(h0), v2 = source(h1), v3 = source(h2)
//   h_alpha[h0] = α3  (angle at v3, opposite edge h0 = v1v2)
//   h_alpha[h1] = α1  (angle at v1, opposite edge h1 = v2v3)
//   h_alpha[h2] = α2  (angle at v2, opposite edge h2 = v3v1)
//
// Property-map name prefix: "ev:" (vertex) and "ee:" (edge).

#include "conformal_mesh.hpp"
#include "constants.hpp"
#include "euclidean_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 Euclidean functional.
using EuclVMapD = ConformalMesh::Property_map<Vertex_index, double>;
/// Property map vertex → `int`    for the Euclidean functional.
using EuclVMapI = ConformalMesh::Property_map<Vertex_index, int>;
/// Property map edge   → `double` for the Euclidean functional.
using EuclEMapD = ConformalMesh::Property_map<Edge_index,   double>;
/// Property map edge   → `int`    for the Euclidean functional.
using EuclEMapI = ConformalMesh::Property_map<Edge_index,   int>;

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

/// Bundle of the five property maps consumed by the Euclidean functional.
struct EuclideanMaps {
    EuclVMapI v_idx;    ///< DOF index per vertex   (-1 = pinned / u_v = 0)
    EuclEMapI e_idx;    ///< DOF index per edge      (-1 = no edge DOF)
    EuclVMapD theta_v;  ///< target cone angle Θ_v   (default 2π)
    EuclEMapD phi_e;    ///< target edge turn angle φ_e (default π)
    EuclEMapD lambda0;  ///< base log-length λ°_e    (default 0.0)
};

/// Attach the five Euclidean property maps to `mesh` with sensible
/// defaults and return their handles.
///
/// Defaults:
/// * `v_idx[v] = -1`  (every vertex pinned; user must reassign before solving)
/// * `e_idx[e] = -1`  (no edge DOFs by default; use `assign_euclidean_all_dof_indices` for cyclic functional)
/// * `theta_v[v] = 2π`  (flat interior vertex target)
/// * `phi_e[e] = π`     (interior edge turn angle target — flat surface)
/// * `lambda0[e] = 0`   (placeholder; call `compute_euclidean_lambda0_from_mesh` next)
///
/// Map name prefix: `"ev:"` (vertex) and `"ee:"` (edge).
inline EuclideanMaps setup_euclidean_maps(ConformalMesh& mesh)
{
    EuclideanMaps m;
    m.v_idx  = mesh.add_property_map<Vertex_index, int>   ("ev:idx",   -1     ).first;
    m.e_idx  = mesh.add_property_map<Edge_index,   int>   ("ee:idx",   -1     ).first;
    m.theta_v= mesh.add_property_map<Vertex_index, double>("ev:theta", TWO_PI ).first;
    m.phi_e  = mesh.add_property_map<Edge_index,   double>("ee:phi",   PI     ).first;
    m.lambda0= mesh.add_property_map<Edge_index,   double>("ee:lam0",  0.0    ).first;
    return m;
}

/// Assign sequential DOF indices `0..n-1` to all vertices.
///
/// **Note:** does NOT pin a gauge vertex.  For closed meshes the caller
/// must set one `m.v_idx[v] = -1` either before or after this call to
/// remove the rotational mode (the Newton solver's SparseQR fallback
/// will otherwise pick a minimum-norm solution but at higher cost).
inline int assign_euclidean_vertex_dof_indices(ConformalMesh& mesh, EuclideanMaps& 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).  Use this overload for the "cyclic" formulation that
/// includes per-edge log-length DOFs (`λ_e`) on top of per-vertex scale
/// factors (`u_v`).
inline int assign_euclidean_all_dof_indices(ConformalMesh& mesh, EuclideanMaps& 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 euclidean_dimension(const ConformalMesh& mesh, const EuclideanMaps& 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;
}

/// Set `lambda0` from mesh vertex positions:
/// `λ°_e = 2·log(|p_i − p_j|)` (natural log of Euclidean edge length²).
/// This gives `exp(Λ̃_ij / 2) = l_ij` at `x = 0`.
inline void compute_euclidean_lambda0_from_mesh(ConformalMesh& mesh, EuclideanMaps& 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));
        double dx = p1.x() - p2.x();
        double dy = p1.y() - p2.y();
        double dz = p1.z() - p2.z();
        double len = std::sqrt(dx*dx + dy*dy + dz*dz);
        if (len > 1e-15)
            m.lambda0[e] = 2.0 * std::log(len);
        else
            m.lambda0[e] = -30.0;   // degenerate edge
    }
}

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

/// Read DOF value from `x` for index `idx`; return 0 if pinned (idx < 0).
static inline double eucl_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 eucl_hidx(Halfedge_index h)
{
    return static_cast<std::size_t>(static_cast<std::uint32_t>(h));
}

/// Compute the Euclidean-functional gradient G(x):
///   * `G_v = Θ_v − Σ_faces α_v(face)`
///   * `G_e = α_opp(face⁺) + α_opp(face⁻) − φ_e`
///
/// Same half-edge corner-angle storage convention as `spherical_gradient`.
inline std::vector<double> euclidean_gradient(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const EuclideanMaps&       m)
{
    const int n = euclidean_dimension(mesh, m);
    std::vector<double> G(static_cast<std::size_t>(n), 0.0);

    // Per-halfedge corner-angle storage.
    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-lengths (always additive: u_i + u_j regardless of DOF status)
        double u1 = eucl_dof_val(m.v_idx[v1], x);
        double u2 = eucl_dof_val(m.v_idx[v2], x);
        double u3 = eucl_dof_val(m.v_idx[v3], x);

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

        auto fa = euclidean_angles(lam12, lam23, lam31);
        if (!fa.valid) continue;   // degenerate triangle: contributes 0

        // h_alpha[h] = corner angle OPPOSITE to h's edge:
        //   h0 (edge v1v2) → opposite corner at v3 → α3
        //   h1 (edge v2v3) → opposite corner at v1 → α1
        //   h2 (edge v3v1) → opposite corner at v2 → α2
        h_alpha[eucl_hidx(h0)] = fa.alpha3;
        h_alpha[eucl_hidx(h1)] = fa.alpha1;
        h_alpha[eucl_hidx(h2)] = fa.alpha2;
    }

    // ── Pass 2: accumulate vertex gradient ───────────────────────────────────
    // 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[eucl_hidx(mesh.prev(h))];
        }
        G[static_cast<std::size_t>(iv)] = m.theta_v[v] - sum_alpha;
    }

    // ── Pass 3: accumulate edge gradient ─────────────────────────────────────
    // G_e = α_opp(f⁺) + α_opp(f⁻) − φ_e
    // α_opp of edge e in face f = h_alpha[halfedge h of e pointing INTO 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 = -m.phi_e[e];
        if (!mesh.is_border(h))  sum += h_alpha[eucl_hidx(h)];
        if (!mesh.is_border(ho)) sum += h_alpha[eucl_hidx(ho)];
        G[static_cast<std::size_t>(ie)] = sum;
    }

    return G;
}

/// Euclidean energy `E(x) = ∫₀¹ ⟨G(t·x), x⟩ dt`, evaluated with
/// 10-point Gauss-Legendre quadrature (same as the Spherical functional).
inline double euclidean_energy(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const EuclideanMaps&       m)
{
    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;
        double wt = gl_w[k] * 0.5;

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

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

        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) ──────────────────────────────────────

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

/// Evaluate the Euclidean functional at DOFs `x`.  Returns energy and
/// gradient (toggle via `need_energy` / `need_gradient`).
inline EuclideanResult evaluate_euclidean(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const EuclideanMaps&       m,
    bool                       need_energy   = true,
    bool                       need_gradient = true)
{
    EuclideanResult res;
    if (need_gradient)
        res.gradient = euclidean_gradient(mesh, x, m);
    if (need_energy)
        res.energy = euclidean_energy(mesh, x, m);
    return res;
}

/// Finite-difference gradient check for the Euclidean functional
/// (central differences).  Defaults `eps = 1e-5`, `tol = 1e-4`.
inline bool gradient_check_euclidean(
    ConformalMesh&             mesh,
    const std::vector<double>& x0,
    const EuclideanMaps&       m,
    double                     eps = 1e-5,
    double                     tol = 1e-4)
{
    auto G = euclidean_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 = euclidean_energy(mesh, xp, m);
        double Em = euclidean_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;
}

} // namespace conformallab