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

// test_newton_solver.cpp
//
// Phase 4 — Newton solver tests.
//
// Design principle:
//   We test convergence to a KNOWN equilibrium.  For the spherical tetrahedron
//   x* = 0 is built-in (G(0) ≈ 0 by construction).  For Euclidean meshes we
//   use "natural theta": set theta_v[v] = actual angle sum at x=0, which makes
//   x* = 0 the exact equilibrium by definition.
//   For HyperIdeal we use the same "natural target" trick at a valid base point
//   (b=1.0, a=0.5), since x=0 is degenerate for the HyperIdeal functional.
//
// Tests:
//   Spherical:
//     1. Converges from x=[-0.2,...] to x*=0 (spherical tetrahedron).
//     2. Converges in few iterations (quadratic convergence near equilibrium).
//     3. Converges from a large perturbation x=[-0.5,...].
//     4. Result fields (x.size, grad_inf_norm) are self-consistent.
//
//   Euclidean:
//     5. Converges (triangle, 1 pinned vertex, natural theta).
//     6. Converges (quad strip, 1 pinned vertex, natural theta).
//     7. Converges with explicitly chosen mixed pinned/variable layout.
//
//   HyperIdeal:
//     8.  Converges on triangle (all variable, natural targets).
//     9.  Result fields self-consistent.
//    10.  Converges on tetrahedron (10 DOFs, larger mesh).
//    11.  SparseQR: result consistent (valid starting region).
//
//   SparseQR fallback (direct unit tests):
//    12.  solve_linear_system recovers correct solution on rank-deficient matrix.
//    13.  solve_linear_system sets fallback_used=true on a singular matrix.
//    14.  Euclidean Newton on closed tetrahedron (no pinned vertex) converges
//         via SparseQR gauge-mode handling.

#include "conformal_mesh.hpp"
#include "mesh_builder.hpp"
#include "euclidean_functional.hpp"
#include "spherical_functional.hpp"
#include "hyper_ideal_functional.hpp"
#include "newton_solver.hpp"
#include <Eigen/Dense>
#include <gtest/gtest.h>
#include <cmath>
#include <vector>

using namespace conformallab;

// ────────────────────────────────────────────────────────────────────────────
// Helper: set theta_v[v] = actual angle sum at x=0 for each variable vertex.
//
// Requires that DOF indices have already been assigned (v_idx populated).
// Uses euclidean_gradient directly so the formula is exact and consistent.
// ────────────────────────────────────────────────────────────────────────────
static void set_natural_euclidean_theta(ConformalMesh& mesh, EuclideanMaps& maps, int n)
{
    std::vector<double> x0(static_cast<std::size_t>(n), 0.0);
    // G[iv] = theta_v[v] - sum_alpha(x=0)
    // so  sum_alpha(x=0) = theta_v[v] - G[iv]
    auto G = euclidean_gradient(mesh, x0, maps);
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv < 0) continue;
        maps.theta_v[v] -= G[static_cast<std::size_t>(iv)];
    }
}

// ════════════════════════════════════════════════════════════════════════════
// Spherical 1 — Converges from moderate perturbation
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Spherical_ConvergesFromPerturbation)
{
    auto mesh = make_spherical_tetrahedron();
    auto maps = setup_spherical_maps(mesh);
    compute_lambda0_from_mesh(mesh, maps);
    int n = assign_vertex_dof_indices(mesh, maps);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.2);
    auto res = newton_spherical(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/50);

    EXPECT_TRUE(res.converged)
        << "Newton (spherical) should converge; grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-8);
}

// ════════════════════════════════════════════════════════════════════════════
// Spherical 2 — Quadratic convergence: few iterations suffice
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Spherical_FewIterations)
{
    auto mesh = make_spherical_tetrahedron();
    auto maps = setup_spherical_maps(mesh);
    compute_lambda0_from_mesh(mesh, maps);
    int n = assign_vertex_dof_indices(mesh, maps);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.2);
    auto res = newton_spherical(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/50);

    EXPECT_LE(res.iterations, 20)
        << "Newton should converge in ≤ 20 iterations; took " << res.iterations;
}

// ════════════════════════════════════════════════════════════════════════════
// Spherical 3 — Large perturbation: global convergence via line search
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Spherical_ConvergesFromLargePerturbation)
{
    auto mesh = make_spherical_tetrahedron();
    auto maps = setup_spherical_maps(mesh);
    compute_lambda0_from_mesh(mesh, maps);
    int n = assign_vertex_dof_indices(mesh, maps);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.5);
    auto res = newton_spherical(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/100);

    EXPECT_TRUE(res.converged)
        << "Newton (spherical, large perturbation) should converge; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-8);
}

// ════════════════════════════════════════════════════════════════════════════
// Spherical 4 — Result fields are self-consistent
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Spherical_ResultFieldsConsistent)
{
    auto mesh = make_spherical_tetrahedron();
    auto maps = setup_spherical_maps(mesh);
    compute_lambda0_from_mesh(mesh, maps);
    int n = assign_vertex_dof_indices(mesh, maps);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.1);
    auto res = newton_spherical(mesh, x0, maps, /*tol=*/1e-8);

    EXPECT_EQ(static_cast<int>(res.x.size()), n);

    // Reported grad_inf_norm must match re-computed gradient at res.x
    auto G = spherical_gradient(mesh, res.x, maps);
    double actual_inf = 0.0;
    for (double v : G) actual_inf = std::max(actual_inf, std::abs(v));
    EXPECT_NEAR(actual_inf, res.grad_inf_norm, 1e-10);
}

// ════════════════════════════════════════════════════════════════════════════
// Euclidean 1 — Triangle with 1 pinned vertex + natural theta
//
// make_triangle(): 3 vertices.  Pin v0 → 2 free DOFs.
// With natural theta: x* = 0 (G(0) = 0 by construction).
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Euclidean_ConvergesTrianglePinned)
{
    auto mesh = make_triangle();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);

    // Pin first vertex, assign sequential DOFs to the other two
    auto vit = mesh.vertices().begin();
    Vertex_index v0 = *vit++;
    maps.v_idx[v0] = -1;     // pinned
    int idx = 0;
    for (; vit != mesh.vertices().end(); ++vit)
        maps.v_idx[*vit] = idx++;
    int n = idx;             // = 2

    set_natural_euclidean_theta(mesh, maps, n);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.1);
    auto res = newton_euclidean(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/50);

    EXPECT_TRUE(res.converged)
        << "Newton (Euclidean, triangle, pinned) should converge; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-8);
}

// ════════════════════════════════════════════════════════════════════════════
// Euclidean 2 — Quad strip with 1 pinned vertex + natural theta
//
// make_quad_strip(): 4 vertices, 2 faces.  Pin v0 → 3 free DOFs.
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Euclidean_ConvergesQuadStripPinned)
{
    auto mesh = make_quad_strip();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);

    // Pin first vertex
    auto vit = mesh.vertices().begin();
    Vertex_index v0 = *vit++;
    maps.v_idx[v0] = -1;
    int idx = 0;
    for (; vit != mesh.vertices().end(); ++vit)
        maps.v_idx[*vit] = idx++;
    int n = idx;             // = 3

    set_natural_euclidean_theta(mesh, maps, n);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.15);
    auto res = newton_euclidean(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/50);

    EXPECT_TRUE(res.converged)
        << "Newton (Euclidean, quad strip, pinned) should converge; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-8);
}

// ════════════════════════════════════════════════════════════════════════════
// Euclidean 3 — Mixed pinned layout: explicit vertex assignment
//
// Quad strip: v0 pinned, v1/v2/v3 free.
// Natural theta set AFTER DOF assignment so that x* = 0 is the equilibrium
// for the free vertices (with v0 fixed at u0=0).
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, Euclidean_ConvergesMixedPinned)
{
    auto mesh = make_quad_strip();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);

    // Explicitly assign DOF indices
    auto vit = mesh.vertices().begin();
    Vertex_index v0 = *vit++;
    Vertex_index v1 = *vit++;
    Vertex_index v2 = *vit++;
    Vertex_index v3 = *vit;

    maps.v_idx[v0] = -1;   // pinned at u0 = 0
    maps.v_idx[v1] =  0;
    maps.v_idx[v2] =  1;
    maps.v_idx[v3] =  2;
    const int n = 3;

    // Set natural theta AFTER pinning so that x* = [0,0,0] is the equilibrium
    set_natural_euclidean_theta(mesh, maps, n);

    std::vector<double> x0 = {-0.1, -0.15, -0.05};
    auto res = newton_euclidean(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/50);

    EXPECT_TRUE(res.converged)
        << "Newton (Euclidean, mixed pinned) should converge; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-8);
}

// ════════════════════════════════════════════════════════════════════════════
// Helper: set HyperIdeal target angles to actual sums at a non-degenerate
// base point (b_base, a_base), making that point the equilibrium x*.
//
// Note: x = 0 is degenerate for the HyperIdeal functional (log-space; the
// functional requires b_i > 0 / a_e > 0).  We therefore choose a valid base
// point, evaluate G there, and absorb G into the targets so that G(xbase) = 0.
// Newton tests then start from a perturbation of xbase.
//
// Returns xbase so callers can construct a perturbed starting point.
// ════════════════════════════════════════════════════════════════════════════
static std::vector<double> set_natural_hyper_ideal_targets(
    ConformalMesh& mesh, HyperIdealMaps& maps, int n,
    double b_base = 1.0, double a_base = 0.5)
{
    const auto sz = static_cast<std::size_t>(n);
    std::vector<double> xbase(sz, 0.0);

    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv >= 0) xbase[static_cast<std::size_t>(iv)] = b_base;
    }
    for (auto e : mesh.edges()) {
        int ie = maps.e_idx[e];
        if (ie >= 0) xbase[static_cast<std::size_t>(ie)] = a_base;
    }

    // G = Σβ - theta_target (initial target = 0 → G = Σβ = "actual" angles)
    auto G = evaluate_hyper_ideal(mesh, xbase, maps, /*energy=*/false).gradient;

    // Set target := actual so that G(xbase) = actual - target = 0
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv < 0) continue;
        maps.theta_v[v] += G[static_cast<std::size_t>(iv)];
    }
    for (auto e : mesh.edges()) {
        int ie = maps.e_idx[e];
        if (ie < 0) continue;
        maps.theta_e[e] += G[static_cast<std::size_t>(ie)];
    }

    return xbase;
}

// ════════════════════════════════════════════════════════════════════════════
// HyperIdeal 1 — Triangle, all DOFs variable, converges from perturbation
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, HyperIdeal_ConvergesTriangleAllVariable)
{
    auto mesh = make_triangle();
    auto maps = setup_hyper_ideal_maps(mesh);
    int n = assign_all_dof_indices(mesh, maps);

    // xbase = (b=1.0, a=0.5) is the equilibrium after natural-target setup.
    auto xbase = set_natural_hyper_ideal_targets(mesh, maps, n);

    // Perturb by +0.2 uniformly
    std::vector<double> x0 = xbase;
    for (auto& v : x0) v += 0.2;

    auto res = newton_hyper_ideal(mesh, x0, maps, /*tol=*/1e-7, /*max_iter=*/100);

    EXPECT_TRUE(res.converged)
        << "Newton (HyperIdeal, triangle) should converge; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-7);
}

// ════════════════════════════════════════════════════════════════════════════
// HyperIdeal 2 — Result fields self-consistent
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, HyperIdeal_ResultFieldsConsistent)
{
    auto mesh = make_triangle();
    auto maps = setup_hyper_ideal_maps(mesh);
    int n = assign_all_dof_indices(mesh, maps);

    auto xbase = set_natural_hyper_ideal_targets(mesh, maps, n);

    std::vector<double> x0 = xbase;
    for (auto& v : x0) v += 0.1;

    auto res = newton_hyper_ideal(mesh, x0, maps, /*tol=*/1e-7, /*max_iter=*/100);

    EXPECT_EQ(static_cast<int>(res.x.size()), n);

    // Reported grad_inf_norm must match re-computed gradient at res.x
    auto G = evaluate_hyper_ideal(mesh, res.x, maps, false).gradient;
    double actual_inf = 0.0;
    for (double v : G) actual_inf = std::max(actual_inf, std::abs(v));
    EXPECT_NEAR(actual_inf, res.grad_inf_norm, 1e-9);
}

// ════════════════════════════════════════════════════════════════════════════
// HyperIdeal 3 — Tetrahedron (10 DOFs): 4 vertex b-vals + 6 edge a-vals
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, HyperIdeal_ConvergesTetrahedron)
{
    auto mesh = make_tetrahedron();
    auto maps = setup_hyper_ideal_maps(mesh);
    int n = assign_all_dof_indices(mesh, maps);

    auto xbase = set_natural_hyper_ideal_targets(mesh, maps, n);

    // Perturb by +0.15
    std::vector<double> x0 = xbase;
    for (auto& v : x0) v += 0.15;

    auto res = newton_hyper_ideal(mesh, x0, maps, /*tol=*/1e-7, /*max_iter=*/200);

    EXPECT_TRUE(res.converged)
        << "Newton (HyperIdeal, tetrahedron) should converge; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-7);
}

// ════════════════════════════════════════════════════════════════════════════
// HyperIdeal 4 — SparseQR fallback: solver returns a result (no crash)
//
// With all targets = 0 the equilibrium is not at x=0 but the solver should
// at minimum not crash and return a consistent result struct.
// ════════════════════════════════════════════════════════════════════════════

TEST(NewtonSolver, HyperIdeal_SparseQRFallbackNoCrash)
{
    auto mesh = make_triangle();
    auto maps = setup_hyper_ideal_maps(mesh);
    int n = assign_all_dof_indices(mesh, maps);

    // Leave targets at their default (0): solver tries to solve but the
    // "equilibrium" is at some unknown x*.  With valid starting point the
    // Hessian is positive-definite and the solver should not crash.
    // We don't assert convergence — just that the result struct is consistent.
    std::vector<double> x0(static_cast<std::size_t>(n), 1.0);
    // Mix vertex / edge DOFs: b=1.0, a=0.5 (valid region of the functional)
    for (auto e : mesh.edges()) {
        int ie = maps.e_idx[e];
        if (ie >= 0) x0[static_cast<std::size_t>(ie)] = 0.5;
    }
    auto res = newton_hyper_ideal(mesh, x0, maps, /*tol=*/1e-7, /*max_iter=*/50);

    // Struct fields must always be populated
    EXPECT_EQ(static_cast<int>(res.x.size()), n);
    EXPECT_GE(res.iterations, 0);
    EXPECT_FALSE(std::isnan(res.grad_inf_norm));
    EXPECT_FALSE(std::isinf(res.grad_inf_norm));
}

// ════════════════════════════════════════════════════════════════════════════
// SparseQR fallback — Test 12: solve_linear_system recovers correct solution
//
// The public API solve_linear_system(A, rhs) must return the correct answer
// for a well-conditioned full-rank system (LDLT path taken).
// ════════════════════════════════════════════════════════════════════════════

TEST(SparseQRFallback, FullRankSystem_CorrectSolution)
{
    // Build a simple 3×3 diagonal PD matrix: A = diag(1, 2, 3)
    Eigen::SparseMatrix<double> A(3, 3);
    A.insert(0, 0) = 1.0;
    A.insert(1, 1) = 2.0;
    A.insert(2, 2) = 3.0;
    A.makeCompressed();

    Eigen::VectorXd rhs(3);
    rhs << 1.0, 4.0, 9.0;   // solution = [1, 2, 3]

    bool fallback = true;     // expect it to be set to false (LDLT succeeds)
    Eigen::VectorXd x = conformallab::solve_linear_system(A, rhs, &fallback);

    EXPECT_FALSE(fallback) << "Full-rank system: LDLT should succeed (no SparseQR needed)";
    EXPECT_NEAR(x[0], 1.0, 1e-12);
    EXPECT_NEAR(x[1], 2.0, 1e-12);
    EXPECT_NEAR(x[2], 3.0, 1e-12);
}

// ════════════════════════════════════════════════════════════════════════════
// SparseQR fallback — Test 13: fallback_used=true on a singular matrix
//
// Construct a symmetric 3×3 matrix of rank 1 where LDLT fails (the (2,2)
// pivot is zero).  SparseQR finds the minimum-norm least-squares solution.
// ════════════════════════════════════════════════════════════════════════════

TEST(SparseQRFallback, SingularMatrix_FallbackActivated)
{
    // A = [[2, 0, 0],
    //      [0, 0, 0],   ← zero pivot → LDLT failure
    //      [0, 0, 3]]
    // rhs compatible with the row space: [2, 0, 3] → solution [1, 0, 1]
    Eigen::SparseMatrix<double> A(3, 3);
    A.insert(0, 0) = 2.0;
    // row/col 1 deliberately all-zero
    A.insert(2, 2) = 3.0;
    A.makeCompressed();

    Eigen::VectorXd rhs(3);
    rhs << 2.0, 0.0, 3.0;

    bool fallback = false;
    Eigen::VectorXd x = conformallab::solve_linear_system(A, rhs, &fallback);

    EXPECT_TRUE(fallback) << "Singular matrix: SparseQR fallback must be triggered";
    // SparseQR min-norm solution: x[0]=1, x[1]=0, x[2]=1
    EXPECT_NEAR(x[0], 1.0, 1e-10);
    EXPECT_NEAR(x[1], 0.0, 1e-10);
    EXPECT_NEAR(x[2], 1.0, 1e-10);
}

// ════════════════════════════════════════════════════════════════════════════
// SparseQR fallback — Test 14: Euclidean Newton on a closed mesh, no pinning
//
// make_tetrahedron() is a closed surface (4 vertices, 4 faces).  Without a
// pinned vertex the Euclidean Hessian has a 1-D null space (uniform scale
// gauge mode): H·1 = 0.  SimplicialLDLT fails on this rank-deficient H;
// SparseQR finds the min-norm Newton step orthogonal to the null space.
//
// The gradient always lives in the row space of H (Σ G_v = 0 by angle-sum
// invariance), so the SparseQR step is also the Newton step and the solver
// converges to the natural equilibrium.
// ════════════════════════════════════════════════════════════════════════════

TEST(SparseQRFallback, Euclidean_ClosedMeshNoPinConverges)
{
    auto mesh = make_tetrahedron();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);

    // Assign all 4 vertices as free DOFs (no pinning).
    int idx = 0;
    for (auto v : mesh.vertices())
        maps.v_idx[v] = idx++;
    const int n = idx;   // = 4

    // Natural theta: equilibrium at x* = 0.
    set_natural_euclidean_theta(mesh, maps, n);

    std::vector<double> x0(static_cast<std::size_t>(n), -0.1);
    auto res = newton_euclidean(mesh, x0, maps, /*tol=*/1e-8, /*max_iter=*/100);

    EXPECT_TRUE(res.converged)
        << "Euclidean Newton on closed tetrahedron (no pin) must converge via SparseQR; "
           "grad_inf_norm = " << res.grad_inf_norm;
    EXPECT_LT(res.grad_inf_norm, 1e-8);
}