// test_newton_solver.cpp
//
// Phase 4a — 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.
//
// 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.

#include "conformal_mesh.hpp"
#include "mesh_builder.hpp"
#include "euclidean_functional.hpp"
#include "spherical_functional.hpp"
#include "newton_solver.hpp"
#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);
}