// test_layout.cpp
//
// Phase 5 — Layout / embedding tests.
//
// Strategy: a conformal map that is the identity (natural equilibrium x*=0)
// should reproduce the mesh's own edge lengths.  We verify:
//   1. Euclidean layout preserves edge lengths (to solver tolerance).
//   2. Spherical layout preserves arc-lengths on S².
//   3. Layout2D has correct size (one entry per vertex).
//   4. Serialisation round-trip: save JSON → load → DOF vector unchanged.
//   5. Serialisation round-trip: save XML  → load → DOF vector unchanged.
//   6. HyperIdeal layout returns success and finite positions.

#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 "layout.hpp"
#include "serialization.hpp"
#include <gtest/gtest.h>
#include <cmath>
#include <vector>
#include <filesystem>

using namespace conformallab;

// ────────────────────────────────────────────────────────────────────────────
// Helpers
// ────────────────────────────────────────────────────────────────────────────

// Euclidean edge length from layout (2D)
static double layout_edge_len(const Layout2D& lay,
                               ConformalMesh& mesh, Halfedge_index h)
{
    auto vi = mesh.source(h).idx();
    auto vj = mesh.target(h).idx();
    return (lay.uv[vi] - lay.uv[vj]).norm();
}

// Spherical arc-length from layout (3D)
static double layout_arc_len(const Layout3D& lay,
                              ConformalMesh& mesh, Halfedge_index h)
{
    auto& a = lay.pos[mesh.source(h).idx()];
    auto& b = lay.pos[mesh.target(h).idx()];
    double c = std::max(-1.0, std::min(1.0, a.dot(b)));
    return std::acos(c);
}

// Euclidean edge length from maps + x (the "true" target)
static double maps_edge_len(const EuclideanMaps& m,
                             ConformalMesh& mesh,
                             Halfedge_index h,
                             const std::vector<double>& x)
{
    auto get_u = [&](Vertex_index v) -> double {
        int iv = m.v_idx[v]; return (iv >= 0) ? x[static_cast<std::size_t>(iv)] : 0.0;
    };
    Edge_index e = mesh.edge(h);
    double lam = m.lambda0[e] + get_u(mesh.source(h)) + get_u(mesh.target(h));
    return std::exp(lam * 0.5);
}

// Spherical arc-length from maps + x
static double maps_arc_len(const SphericalMaps& m,
                            ConformalMesh& mesh,
                            Halfedge_index h,
                            const std::vector<double>& x)
{
    auto get_u = [&](Vertex_index v) -> double {
        int iv = m.v_idx[v]; return (iv >= 0) ? x[static_cast<std::size_t>(iv)] : 0.0;
    };
    Edge_index e = mesh.edge(h);
    double lam = m.lambda0[e] + get_u(mesh.source(h)) + get_u(mesh.target(h));
    return 2.0 * std::asin(std::min(std::exp(lam * 0.5), 1.0));
}

// ════════════════════════════════════════════════════════════════════════════
// Test 1 — Euclidean layout preserves edge lengths (identity map)
// ════════════════════════════════════════════════════════════════════════════

TEST(Layout, Euclidean_PreservesEdgeLengths)
{
    auto mesh = make_quad_strip();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);

    // Pin first vertex; natural equilibrium so x* = 0
    auto vit = mesh.vertices().begin();
    maps.v_idx[*vit++] = -1;
    int idx = 0;
    for (; vit != mesh.vertices().end(); ++vit) maps.v_idx[*vit] = idx++;
    const int n = idx;

    std::vector<double> x0(static_cast<std::size_t>(n), 0.0);
    auto G0 = euclidean_gradient(mesh, x0, maps);
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv >= 0) maps.theta_v[v] -= G0[static_cast<std::size_t>(iv)];
    }

    // Solve (identity map)
    auto res = newton_euclidean(mesh, x0, maps, 1e-10, 100);
    ASSERT_TRUE(res.converged);

    auto layout = euclidean_layout(mesh, res.x, maps);
    ASSERT_TRUE(layout.success);
    EXPECT_EQ(layout.uv.size(), mesh.number_of_vertices());

    // Every edge: |layout_len - expected_len| < 1e-8
    for (auto e : mesh.edges()) {
        Halfedge_index h = mesh.halfedge(e, 0);
        double expected = maps_edge_len(maps, mesh, h, res.x);
        double actual   = layout_edge_len(layout, mesh, h);
        EXPECT_NEAR(actual, expected, 1e-7)
            << "Edge " << e << ": expected " << expected << " got " << actual;
    }
}

// ════════════════════════════════════════════════════════════════════════════
// Test 2 — Euclidean layout: correct vertex count
// ════════════════════════════════════════════════════════════════════════════

TEST(Layout, Euclidean_CorrectVertexCount)
{
    auto mesh = make_triangle();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);

    auto vit = mesh.vertices().begin();
    maps.v_idx[*vit++] = -1;
    int idx = 0;
    for (; vit != mesh.vertices().end(); ++vit) maps.v_idx[*vit] = idx++;
    const int n = idx;

    std::vector<double> x(static_cast<std::size_t>(n), 0.0);
    auto layout = euclidean_layout(mesh, x, maps);

    EXPECT_TRUE(layout.success);
    EXPECT_EQ(static_cast<int>(layout.uv.size()),
              static_cast<int>(mesh.number_of_vertices()));
}

// ════════════════════════════════════════════════════════════════════════════
// Test 3 — Euclidean triangle layout: root face is a valid triangle
// ════════════════════════════════════════════════════════════════════════════

TEST(Layout, Euclidean_TriangleIsNonDegenerate)
{
    auto mesh = make_triangle();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);
    for (auto v : mesh.vertices()) maps.v_idx[v] = static_cast<int>(v.idx());

    std::vector<double> x(mesh.number_of_vertices(), 0.0);
    auto layout = euclidean_layout(mesh, x, maps);
    ASSERT_TRUE(layout.success);

    // Area of layout triangle > 0
    const auto& A = layout.uv[0];
    const auto& B = layout.uv[1];
    const auto& C = layout.uv[2];
    double area = std::abs((B - A).x() * (C - A).y()
                         - (C - A).x() * (B - A).y()) * 0.5;
    EXPECT_GT(area, 1e-10) << "Layout triangle must be non-degenerate";
}

// ════════════════════════════════════════════════════════════════════════════
// Test 4 — Spherical layout: arc-lengths preserved (identity map)
// ════════════════════════════════════════════════════════════════════════════

TEST(Layout, Spherical_PreservesArcLengths)
{
    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);

    // Solve to identity
    std::vector<double> x0(static_cast<std::size_t>(n), 0.0);
    auto G0 = spherical_gradient(mesh, x0, maps);
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv >= 0) maps.theta_v[v] -= G0[static_cast<std::size_t>(iv)];
    }
    auto res = newton_spherical(mesh, x0, maps, 1e-10, 100);
    ASSERT_TRUE(res.converged);

    auto layout = spherical_layout(mesh, res.x, maps);
    ASSERT_TRUE(layout.success);
    EXPECT_EQ(layout.pos.size(), mesh.number_of_vertices());

    for (auto e : mesh.edges()) {
        Halfedge_index h = mesh.halfedge(e, 0);
        double expected = maps_arc_len(maps, mesh, h, res.x);
        double actual   = layout_arc_len(layout, mesh, h);
        EXPECT_NEAR(actual, expected, 1e-6)
            << "Spherical edge " << e << ": expected " << expected << " got " << actual;
    }
}

// ════════════════════════════════════════════════════════════════════════════
// Test 5 — Spherical layout: all positions on unit sphere
// ════════════════════════════════════════════════════════════════════════════

TEST(Layout, Spherical_PositionsOnUnitSphere)
{
    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> x(static_cast<std::size_t>(n), 0.0);
    auto layout = spherical_layout(mesh, x, maps);
    ASSERT_TRUE(layout.success);

    for (auto v : mesh.vertices()) {
        double r = layout.pos[v.idx()].norm();
        EXPECT_NEAR(r, 1.0, 1e-12) << "Vertex " << v << " not on unit sphere, |p| = " << r;
    }
}

// ════════════════════════════════════════════════════════════════════════════
// Test 6 — HyperIdeal layout returns success and finite positions
// ════════════════════════════════════════════════════════════════════════════

TEST(Layout, HyperIdeal_SuccessAndFinitePositions)
{
    auto mesh = make_triangle();
    auto maps = setup_hyper_ideal_maps(mesh);
    int n = assign_all_dof_indices(mesh, maps);

    // Natural equilibrium at (b=1, a=0.5)
    std::vector<double> xbase(static_cast<std::size_t>(n));
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v]; if (iv >= 0) xbase[iv] = 1.0;
    }
    for (auto e : mesh.edges()) {
        int ie = maps.e_idx[e]; if (ie >= 0) xbase[ie] = 0.5;
    }
    auto G0 = evaluate_hyper_ideal(mesh, xbase, maps, false).gradient;
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v]; if (iv >= 0) maps.theta_v[v] += G0[iv];
    }
    for (auto e : mesh.edges()) {
        int ie = maps.e_idx[e]; if (ie >= 0) maps.theta_e[e] += G0[ie];
    }

    auto res = newton_hyper_ideal(mesh, xbase, maps, 1e-9, 100);
    ASSERT_TRUE(res.converged);

    auto layout = hyper_ideal_layout(mesh, res.x, maps);
    EXPECT_TRUE(layout.success);
    ASSERT_EQ(layout.uv.size(), mesh.number_of_vertices());

    for (auto v : mesh.vertices()) {
        auto& p = layout.uv[v.idx()];
        EXPECT_FALSE(std::isnan(p.x())) << "NaN in layout x for v" << v;
        EXPECT_FALSE(std::isnan(p.y())) << "NaN in layout y for v" << v;
        // Poincaré disk: all points within the unit disk
        EXPECT_LE(p.norm(), 1.0 + 1e-9) << "Point outside Poincaré disk for v" << v;
    }
}

// ════════════════════════════════════════════════════════════════════════════
// Test 7 — JSON serialisation round-trip
// ════════════════════════════════════════════════════════════════════════════

TEST(Serialization, JSON_RoundTrip)
{
    auto mesh = make_triangle();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);
    auto vit = mesh.vertices().begin();
    maps.v_idx[*vit++] = -1;
    int idx = 0;
    for (; vit != mesh.vertices().end(); ++vit) maps.v_idx[*vit] = idx++;
    const int n = idx;

    std::vector<double> x0(static_cast<std::size_t>(n), -0.05);
    auto G0 = euclidean_gradient(mesh, std::vector<double>(n, 0.0), maps);
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv >= 0) maps.theta_v[v] -= G0[static_cast<std::size_t>(iv)];
    }
    auto res = newton_euclidean(mesh, x0, maps, 1e-10, 100);
    ASSERT_TRUE(res.converged);

    auto layout = euclidean_layout(mesh, res.x, maps);

    const std::string path = "/tmp/conformallab_test_round_trip.json";
    ASSERT_NO_THROW(save_result_json(path, res, "euclidean",
        static_cast<int>(mesh.number_of_vertices()),
        static_cast<int>(mesh.number_of_faces()),
        &layout));
    ASSERT_TRUE(std::filesystem::exists(path));

    NewtonResult res2;
    std::string geom;
    Layout2D layout2;
    auto x2 = load_result_json(path, &res2, &geom, &layout2);

    EXPECT_EQ(geom, "euclidean");
    EXPECT_EQ(res2.converged,  res.converged);
    EXPECT_EQ(res2.iterations, res.iterations);
    ASSERT_EQ(x2.size(), res.x.size());
    for (std::size_t i = 0; i < x2.size(); ++i)
        EXPECT_NEAR(x2[i], res.x[i], 1e-14);

    EXPECT_TRUE(layout2.success);
    ASSERT_EQ(layout2.uv.size(), layout.uv.size());
    for (std::size_t i = 0; i < layout2.uv.size(); ++i) {
        EXPECT_NEAR(layout2.uv[i].x(), layout.uv[i].x(), 1e-12);
        EXPECT_NEAR(layout2.uv[i].y(), layout.uv[i].y(), 1e-12);
    }
    std::filesystem::remove(path);
}

// ════════════════════════════════════════════════════════════════════════════
// Test 8 — XML serialisation round-trip
// ════════════════════════════════════════════════════════════════════════════

TEST(Serialization, XML_RoundTrip)
{
    auto mesh = make_triangle();
    auto maps = setup_euclidean_maps(mesh);
    compute_euclidean_lambda0_from_mesh(mesh, maps);
    auto vit = mesh.vertices().begin();
    maps.v_idx[*vit++] = -1;
    int idx = 0;
    for (; vit != mesh.vertices().end(); ++vit) maps.v_idx[*vit] = idx++;
    const int n = idx;

    std::vector<double> x0(static_cast<std::size_t>(n), 0.0);
    auto G0 = euclidean_gradient(mesh, x0, maps);
    for (auto v : mesh.vertices()) {
        int iv = maps.v_idx[v];
        if (iv >= 0) maps.theta_v[v] -= G0[static_cast<std::size_t>(iv)];
    }
    auto res = newton_euclidean(mesh, x0, maps, 1e-10, 100);
    ASSERT_TRUE(res.converged);

    auto layout = euclidean_layout(mesh, res.x, maps);

    const std::string path = "/tmp/conformallab_test_round_trip.xml";
    ASSERT_NO_THROW(save_result_xml(path, res, "euclidean",
        static_cast<int>(mesh.number_of_vertices()),
        static_cast<int>(mesh.number_of_faces()),
        &layout));
    ASSERT_TRUE(std::filesystem::exists(path));

    NewtonResult res2;
    std::string geom;
    Layout2D layout2;
    auto x2 = load_result_xml(path, &res2, &geom, &layout2);

    EXPECT_EQ(geom, "euclidean");
    EXPECT_EQ(res2.converged,  res.converged);
    ASSERT_EQ(x2.size(), res.x.size());
    for (std::size_t i = 0; i < x2.size(); ++i)
        EXPECT_NEAR(x2[i], res.x[i], 1e-12);

    EXPECT_TRUE(layout2.success);
    ASSERT_EQ(layout2.uv.size(), layout.uv.size());

    std::filesystem::remove(path);
}
