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

// hyper_ideal_functional.hpp
//
// Energy and gradient of the hyper-ideal discrete conformal map functional
// evaluated on a ConformalMesh (CGAL::Surface_mesh).
//
// Ported from de.varylab.discreteconformal.functional.HyperIdealFunctional.
//
// ┌─────────────────────────────────────────────────────────────────────────┐
// │  E(x) = Σ_faces U(f)  −  Σ_edges θ_e · a_e  −  Σ_vertices Θ_v · b_v  │
// │                                                                         │
// │  ∂E/∂b_v = Σ_{faces adj. v} β_v(face)  −  Θ_v                         │
// │  ∂E/∂a_e = α_e(face⁺) + α_e(face⁻)    −  θ_e                         │
// └─────────────────────────────────────────────────────────────────────────┘
//
// DOF vector layout (matches getDimension() ordering):
//   x[v_idx[v]]  = b_v  for each variable vertex   (log scale factor)
//   x[e_idx[e]]  = a_e  for each variable edge      (intersection angle)
//   -1 in v_idx / e_idx means "pinned" (ideal / fixed at 0).
//
// Usage
// ─────
//   auto mesh = make_tetrahedron();
//   auto maps = setup_hyper_ideal_maps(mesh);
//   int n = assign_all_dof_indices(mesh, maps);   // all variable
//   std::vector<double> x(n, 1.0);
//   auto res = evaluate_hyper_ideal(mesh, x, maps);
//   // res.energy, res.gradient

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

namespace conformallab {

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

/// Property map vertex → `double` (HyperIdeal scalar-per-vertex data).
using VMapD = ConformalMesh::Property_map<Vertex_index,   double>;
/// Property map vertex → `int`    (HyperIdeal DOF indices).
using VMapI = ConformalMesh::Property_map<Vertex_index,   int>;
/// Property map edge   → `double` (HyperIdeal scalar-per-edge data).
using EMapD = ConformalMesh::Property_map<Edge_index,     double>;
/// Property map edge   → `int`    (HyperIdeal DOF indices).
using EMapI = ConformalMesh::Property_map<Edge_index,     int>;

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

/// Bundle of the four property maps consumed by the HyperIdeal functional.
struct HyperIdealMaps {
    VMapI v_idx;    ///< DOF index per vertex (−1 = pinned / ideal point).
    EMapI e_idx;    ///< DOF index per edge   (−1 = fixed).
    VMapD theta_v;  ///< Target cone angle Θᵥ (parameter, not variable).
    EMapD theta_e;  ///< Target intersection angle θₑ.
};

/// Attach the four HyperIdeal property maps to `mesh` and return their
/// handles.
///
/// Defaults:
/// * `v_idx[v] = -1`  (ideal vertex — i.e. the corresponding `b_v` is fixed at 0)
/// * `e_idx[e] = -1`  (edge DOF fixed at 0)
/// * `theta_v[v] = 2π`  (regular cone target)
/// * `theta_e[e] = π`   (orthogonal-circle target)
///
/// The map prefix `"v:"` / `"e:"` is intentionally generic for the
/// HyperIdeal functional — it is the canonical / Phase 3b model.
/// Other functionals use distinct prefixes (`"ev:"` Euclidean, `"sv:"`
/// Spherical, `"cf:"`/`"ce:"` CP-Euclidean, `"iv:"`/`"ie:"`
/// Inversive-Distance) so all models can coexist on the same mesh.
inline HyperIdealMaps setup_hyper_ideal_maps(ConformalMesh& mesh)
{
    HyperIdealMaps m;
    m.v_idx   = mesh.add_property_map<Vertex_index, int>   ("v:idx",   -1      ).first;
    m.e_idx   = mesh.add_property_map<Edge_index,   int>   ("e:idx",   -1      ).first;
    m.theta_v = mesh.add_property_map<Vertex_index, double>("v:theta", 2.0*PI  ).first;
    m.theta_e = mesh.add_property_map<Edge_index,   double>("e:theta", PI      ).first;
    return m;
}

/// Count free DOFs: `#variable_vertices + #variable_edges`.
inline int hyper_ideal_dimension(const ConformalMesh& mesh, const HyperIdealMaps& 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;
}

/// Make every vertex hyper-ideal and every edge variable, assigning
/// sequential DOF indices `0..n-1` (vertices first, edges after).
///
/// This is the standard initialisation for the Springborn-2020
/// hyper-ideal functional — gauge fixing is **not** needed because
/// the energy is strictly convex on the full DOF space (no
/// rotational mode for an all-hyper-ideal configuration).
///
/// \returns total DOF count = `num_vertices(mesh) + num_edges(mesh)`.
inline int assign_all_dof_indices(ConformalMesh& mesh, HyperIdealMaps& 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;
}

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

/// Output of `evaluate_hyper_ideal()` — the energy value and (optionally)
/// its gradient evaluated at the current DOF vector.
struct HyperIdealResult {
    double              energy   = 0.0; ///< Functional value at the input DOFs.
    std::vector<double> gradient;       ///< Gradient ∇E; empty when not requested.
};

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

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

// ── Pure-math face-angle kernel ──────────────────────────────────────────────
//
// Computes the six per-face angle outputs (β₁, β₂, β₃, α₁₂, α₂₃, α₃₁) from
// the six local DOF inputs (b₁, b₂, b₃, a₁₂, a₂₃, a₃₁) and the variability
// flags (vᵢb).  This is the pure functional core of `compute_face_angles`
// — no mesh, no property maps, no global x vector.
//
// Why exposed as a free function (Phase 9b):
// ─────────────────────────────────────────
// The block-FD Hessian (`hyper_ideal_hessian_block_fd`) perturbs only the
// 6 DOFs adjacent to a single face at a time, recomputes the 6 angle
// outputs of that face, and uses the local 6×6 Jacobian to scatter into
// the global Hessian.  Working through a pure 6→6 function (instead of
// perturbing the full x and re-running the gradient over all faces)
// reduces the cost of the Hessian from O(F·n) to O(F·36).
//
// The clamping logic (negative b → 0.01, negative a → 0) mirrors
// HyperIdealFunctional.java's defensive behaviour (lines 122-127 of the
// Java original); this keeps the FD perturbation regime well-defined.

/// Six per-face angle outputs computed from local DOFs (see
/// `face_angles_from_local_dofs`).  Used by the block-FD Hessian.
struct FaceAngleOutputs {
    double beta1;   ///< Interior angle at v₁.
    double beta2;   ///< Interior angle at v₂.
    double beta3;   ///< Interior angle at v₃.
    double alpha12; ///< Dihedral angle at edge e₁₂.
    double alpha23; ///< Dihedral angle at edge e₂₃.
    double alpha31; ///< Dihedral angle at edge e₃₁.
};

/// Pure-math 6→6 kernel: given the six local DOFs (b₁,b₂,b₃,a₁₂,a₂₃,a₃₁)
/// of one face plus the per-vertex variability flags, return the six
/// HyperIdeal angle outputs.  No mesh, no property maps — used by the
/// per-face block-FD Hessian in `hyper_ideal_hessian.hpp`.
inline FaceAngleOutputs face_angles_from_local_dofs(
    double b1, double b2, double b3,
    double a12, double a23, double a31,
    bool   v1b, bool v2b, bool v3b)
{
    // Same defensive clamps as compute_face_angles.
    if (v1b && v2b && a12 < 0.0) a12 = 0.0;
    if (v2b && v3b && a23 < 0.0) a23 = 0.0;
    if (v3b && v1b && a31 < 0.0) a31 = 0.0;
    if (v1b && b1  < 0.0) b1 = 0.01;
    if (v2b && b2  < 0.0) b2 = 0.01;
    if (v3b && b3  < 0.0) b3 = 0.01;

    double l12 = lij(b1, b2, a12, v1b, v2b);
    double l23 = lij(b2, b3, a23, v2b, v3b);
    double l31 = lij(b3, b1, a31, v3b, v1b);

    if (l12 < 1E-12 && l23 < 1E-12 && l31 < 1E-12)
        l12 = l23 = l31 = 1E-12;

    FaceAngleOutputs o;

    if (l12 > l23 + l31) {
        o.beta1 = 0.0; o.beta2 = 0.0; o.beta3 = PI;
        o.alpha12 = PI; o.alpha23 = 0.0; o.alpha31 = 0.0;
    } else if (l23 > l12 + l31) {
        o.beta1 = PI; o.beta2 = 0.0; o.beta3 = 0.0;
        o.alpha12 = 0.0; o.alpha23 = PI; o.alpha31 = 0.0;
    } else if (l31 > l12 + l23) {
        o.beta1 = 0.0; o.beta2 = PI; o.beta3 = 0.0;
        o.alpha12 = 0.0; o.alpha23 = 0.0; o.alpha31 = PI;
    } else {
        o.beta1 = zeta(l12, l31, l23);
        o.beta2 = zeta(l23, l12, l31);
        o.beta3 = zeta(l31, l23, l12);
        o.alpha12 = alpha_ij(a12, a23, a31, b1, b2, b3,
                             o.beta1, o.beta2, o.beta3, v1b, v2b, v3b);
        o.alpha23 = alpha_ij(a23, a31, a12, b2, b3, b1,
                             o.beta2, o.beta3, o.beta1, v2b, v3b, v1b);
        o.alpha31 = alpha_ij(a31, a12, a23, b3, b1, b2,
                             o.beta3, o.beta1, o.beta2, v3b, v1b, v2b);
    }
    return o;
}

// ── Per-face angle kernel ─────────────────────────────────────────────────────

/// Per-face angle bundle returned by `compute_face_angles()`.  Carries
/// the six output angles plus the six input DOFs (so the energy and
/// gradient kernels can reuse them without re-reading the mesh).
struct FaceAngles {
    double alpha12; ///< Dihedral angle at edge e₁₂.
    double alpha23; ///< Dihedral angle at edge e₂₃.
    double alpha31; ///< Dihedral angle at edge e₃₁.
    double beta1;   ///< Interior angle at vertex v₁.
    double beta2;   ///< Interior angle at vertex v₂.
    double beta3;   ///< Interior angle at vertex v₃.
    double a12;     ///< Edge DOF value at e₁₂.
    double a23;     ///< Edge DOF value at e₂₃.
    double a31;     ///< Edge DOF value at e₃₁.
    double b1;      ///< Vertex DOF value at v₁.
    double b2;      ///< Vertex DOF value at v₂.
    double b3;      ///< Vertex DOF value at v₃.
    bool   v1b;     ///< `true` iff vertex v₁ is variable (not pinned).
    bool   v2b;     ///< `true` iff vertex v₂ is variable.
    bool   v3b;     ///< `true` iff vertex v₃ is variable.
};

/// Compute the six per-face angles (+ remember the input DOFs) for face
/// `f` of `mesh`, given the current DOF vector `x` and DOF-index maps.
static FaceAngles compute_face_angles(
    const ConformalMesh&       mesh,
    Face_index                 f,
    const std::vector<double>& x,
    const HyperIdealMaps&      m)
{
    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);

    FaceAngles fa;
    fa.v1b = m.v_idx[v1] >= 0;
    fa.v2b = m.v_idx[v2] >= 0;
    fa.v3b = m.v_idx[v3] >= 0;

    fa.a12 = dof_val(m.e_idx[e12], x);
    fa.a23 = dof_val(m.e_idx[e23], x);
    fa.a31 = dof_val(m.e_idx[e31], x);
    fa.b1  = dof_val(m.v_idx[v1],  x);
    fa.b2  = dof_val(m.v_idx[v2],  x);
    fa.b3  = dof_val(m.v_idx[v3],  x);

    // Clamp invalid inputs (mirrors Java log.warning + clamp)
    if (fa.v1b && fa.v2b && fa.a12 < 0.0) fa.a12 = 0.0;
    if (fa.v2b && fa.v3b && fa.a23 < 0.0) fa.a23 = 0.0;
    if (fa.v3b && fa.v1b && fa.a31 < 0.0) fa.a31 = 0.0;
    if (fa.v1b && fa.b1  < 0.0)           fa.b1  = 0.01;
    if (fa.v2b && fa.b2  < 0.0)           fa.b2  = 0.01;
    if (fa.v3b && fa.b3  < 0.0)           fa.b3  = 0.01;

    double l12 = lij(fa.b1, fa.b2, fa.a12, fa.v1b, fa.v2b);
    double l23 = lij(fa.b2, fa.b3, fa.a23, fa.v2b, fa.v3b);
    double l31 = lij(fa.b3, fa.b1, fa.a31, fa.v3b, fa.v1b);

    // Guard degenerate lengths
    if (l12 < 1E-12 && l23 < 1E-12 && l31 < 1E-12)
        l12 = l23 = l31 = 1E-12;

    // Check triangle inequalities; degenerate cases get extreme angles
    if (l12 > l23 + l31) {
        fa.beta1 = 0.0; fa.beta2 = 0.0; fa.beta3 = PI;
        fa.alpha12 = PI; fa.alpha23 = 0.0; fa.alpha31 = 0.0;
    } else if (l23 > l12 + l31) {
        fa.beta1 = PI; fa.beta2 = 0.0; fa.beta3 = 0.0;
        fa.alpha12 = 0.0; fa.alpha23 = PI; fa.alpha31 = 0.0;
    } else if (l31 > l12 + l23) {
        fa.beta1 = 0.0; fa.beta2 = PI; fa.beta3 = 0.0;
        fa.alpha12 = 0.0; fa.alpha23 = 0.0; fa.alpha31 = PI;
    } else {
        fa.beta1 = zeta(l12, l31, l23);
        fa.beta2 = zeta(l23, l12, l31);
        fa.beta3 = zeta(l31, l23, l12);

        fa.alpha12 = alpha_ij(fa.a12, fa.a23, fa.a31,
                              fa.b1,  fa.b2,  fa.b3,
                              fa.beta1, fa.beta2, fa.beta3,
                              fa.v1b, fa.v2b, fa.v3b);
        fa.alpha23 = alpha_ij(fa.a23, fa.a31, fa.a12,
                              fa.b2,  fa.b3,  fa.b1,
                              fa.beta2, fa.beta3, fa.beta1,
                              fa.v2b, fa.v3b, fa.v1b);
        fa.alpha31 = alpha_ij(fa.a31, fa.a12, fa.a23,
                              fa.b3,  fa.b1,  fa.b2,
                              fa.beta3, fa.beta1, fa.beta2,
                              fa.v3b, fa.v1b, fa.v2b);
    }
    return fa;
}

/// Per-face energy contribution U(f) before subtracting the θ·a and Θ·b terms.
static double face_energy(const FaceAngles& fa)
{
    double aa = fa.a12*fa.alpha12 + fa.a23*fa.alpha23 + fa.a31*fa.alpha31;
    double bb = fa.b1 *fa.beta1   + fa.b2 *fa.beta2   + fa.b3 *fa.beta3;

    double V = 0.0;
    if (fa.v1b && fa.v2b && fa.v3b) {
        V = calculateTetrahedronVolume(
                fa.beta1, fa.beta2, fa.beta3,
                fa.alpha23, fa.alpha31, fa.alpha12);
    } else if (!fa.v1b) {
        V = calculateTetrahedronVolumeWithIdealVertexAtGamma(
                fa.beta1, fa.alpha31, fa.alpha12,
                fa.alpha23, fa.beta2, fa.beta3);
    } else if (!fa.v2b) {
        V = calculateTetrahedronVolumeWithIdealVertexAtGamma(
                fa.beta2, fa.alpha12, fa.alpha23,
                fa.alpha31, fa.beta3, fa.beta1);
    } else {   // !v3b
        V = calculateTetrahedronVolumeWithIdealVertexAtGamma(
                fa.beta3, fa.alpha23, fa.alpha31,
                fa.alpha12, fa.beta1, fa.beta2);
    }
    return aa + bb + 2.0 * V;
}

// ── Full evaluation ───────────────────────────────────────────────────────────

/// Evaluate the HyperIdeal functional at DOF vector `x`.  Returns the
/// energy value and (optionally) the gradient in a `HyperIdealResult`.
///
/// \param mesh           Triangle mesh carrying the DOF-index property maps.
/// \param x              Current DOF vector (length = `hyper_ideal_dimension(...)`).
/// \param m              Property-map bundle from `setup_hyper_ideal_maps(...)`.
/// \param need_energy    If `true`, fill `result.energy` (default: `true`).
/// \param need_gradient  If `true`, fill `result.gradient` (default: `true`).
inline HyperIdealResult evaluate_hyper_ideal(
    ConformalMesh&             mesh,
    const std::vector<double>& x,
    const HyperIdealMaps&      m,
    bool                       need_energy   = true,
    bool                       need_gradient = true)
{
    HyperIdealResult res;

    // Temporary per-halfedge storage for computed angles.
    // Indexed by the integer value of Halfedge_index.
    const std::size_t nh = mesh.number_of_halfedges();
    std::vector<double> h_alpha(nh, 0.0);  // α_ij stored on halfedge
    std::vector<double> h_beta (nh, 0.0);  // β_i  stored on opposite halfedge

    // ── Pass 1: angles + energy 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);

        FaceAngles fa = compute_face_angles(mesh, f, x, m);

        // Store computed angles into temporary arrays.
        // h_alpha[h] = α for the edge of h in this face.
        h_alpha[hidx(h0)] = fa.alpha12;
        h_alpha[hidx(h1)] = fa.alpha23;
        h_alpha[hidx(h2)] = fa.alpha31;

        // h_beta[h] = β at the vertex OPPOSITE to h.
        // β1 (at v1 = source(h0)) is opposite to h1 = e23.
        h_beta[hidx(h1)] = fa.beta1;   // h1 is across from v1
        h_beta[hidx(h2)] = fa.beta2;   // h2 is across from v2
        h_beta[hidx(h0)] = fa.beta3;   // h0 is across from v3

        if (need_energy)
            res.energy += face_energy(fa);
    }

    // ── Pass 2: linear energy terms ──────────────────────────────────────────
    if (need_energy) {
        for (auto e : mesh.edges()) {
            int ie = m.e_idx[e];
            if (ie >= 0) res.energy -= m.theta_e[e] * x[static_cast<std::size_t>(ie)];
        }
        for (auto v : mesh.vertices()) {
            int iv = m.v_idx[v];
            if (iv >= 0) res.energy -= m.theta_v[v] * x[static_cast<std::size_t>(iv)];
        }
    }

    // ── Pass 3: gradient ─────────────────────────────────────────────────────
    if (need_gradient) {
        const int n = hyper_ideal_dimension(mesh, m);
        res.gradient.assign(static_cast<std::size_t>(n), 0.0);

        // ∂E/∂b_v = Σ_{faces adj. v} β_v(face)  −  Θ_v
        // β_v(face) = h_beta[prev(h)] for the incoming halfedge h to v in that face.
        for (auto v : mesh.vertices()) {
            int iv = m.v_idx[v];
            if (iv < 0) continue;
            for (auto h : CGAL::halfedges_around_target(v, mesh)) {
                if (mesh.is_border(h)) continue;
                res.gradient[static_cast<std::size_t>(iv)] += h_beta[hidx(mesh.prev(h))];
            }
            res.gradient[static_cast<std::size_t>(iv)] -= m.theta_v[v];
        }

        // ∂E/∂a_e = α_e(face⁺) + α_e(face⁻)  −  θ_e
        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);
            if (!mesh.is_border(h))  res.gradient[static_cast<std::size_t>(ie)] += h_alpha[hidx(h)];
            if (!mesh.is_border(ho)) res.gradient[static_cast<std::size_t>(ie)] += h_alpha[hidx(ho)];
            res.gradient[static_cast<std::size_t>(ie)] -= m.theta_e[e];
        }
    }

    return res;
}

/// Finite-difference gradient check (central differences).
///
/// Returns `true` iff `|G[i] − fd[i]| / max(1, |G[i]|) < tol` for every
/// DOF.  Defaults `eps = 1e-5`, `tol = 1e-4` match the Java `FunctionalTest`.
inline bool gradient_check(
    ConformalMesh&             mesh,
    const std::vector<double>& x0,
    const HyperIdealMaps&      m,
    double                     eps = 1E-5,
    double                     tol = 1E-4)
{
    // Analytic gradient
    auto res = evaluate_hyper_ideal(mesh, x0, m, false, true);
    const auto& G = res.gradient;
    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 = evaluate_hyper_ideal(mesh, xp, m, true, false).energy;
        double Em = evaluate_hyper_ideal(mesh, xm, m, true, false).energy;

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

        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