#include <algorithm>
#include <cmath>
#include <memory>
#include <numbers>
#include <utility>
#include <vector>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/bilinearform.hpp"
#include "mfem/fem/bilininteg.hpp"
#include "mfem/fem/coefficient.hpp"
#include "mfem/fem/gridfunc.hpp"
#include "serif/backend/mfem/shared/boundary_elements.hpp"
#include "serif/backend/mfem/geometry/deformation/reference_geometry.hpp"
#include "serif/backend/mfem/geometry/deformation/surface_metric.hpp"
#include "serif/geometry/deformation/composition.hpp"
#include "serif/geometry/deformation/linearization.hpp"
#include "serif/tests/deformation/mesh.hpp"

namespace {
    namespace deformation = serif::geometry::deformation;
    namespace backend = deformation::backend::mfem;
    using PairedMesh = serif::tests::deformation::PairedMesh;

    /** @brief Independently integrate an analytic monopole or axial-dipole vector trace using MeanField's declared mass quadrature. */
    [[nodiscard]] double analytic_trace_norm(PairedMesh& fixture, int degree) {
        /// Project analytic Cartesian data at FE nodes; evaluating it directly at quadrature points would change the realized trace.
        ::mfem::VectorFunctionCoefficient coefficient(3, [degree](const ::mfem::Vector& position, ::mfem::Vector& value) {
            const double radius = position.Norml2();
            value.SetSize(3);
            value = 0.0;
            if (radius <= 0.0) {
                return;
            }
            const double harmonic = degree == 0 ? 1.0 / std::sqrt(4.0 * std::numbers::pi) : std::sqrt(3.0 / (4.0 * std::numbers::pi)) * position[2] / radius;
            for (int component = 0; component < 3; ++component) {
                value[component] = harmonic * position[component] / radius;
            }
        });
        ::mfem::GridFunction tangent(&fixture.vector_space);
        tangent.ProjectCoefficient(coefficient);
        ::mfem::Array<int> dofs;
        ::mfem::Vector local;
        ::mfem::Vector shape;
        long double work = 0.0L;
        for (int boundary = 0; boundary < fixture.generated.mesh->GetNBE(); ++boundary) {
            if (fixture.generated.mesh->GetBdrAttribute(boundary) != fixture.surface_attribute) {
                continue;
            }
            const ::mfem::FiniteElement& element = *fixture.vector_space.GetBE(boundary);
            ::mfem::ElementTransformation& transformation = *fixture.vector_space.GetBdrElementTransformation(boundary);
            REQUIRE(fixture.vector_space.GetBdrElementVDofs(boundary, dofs) == nullptr);
            tangent.GetSubVector(dofs, local);
            REQUIRE(local.Size() == 3 * element.GetDof());
            shape.SetSize(element.GetDof());
            const ::mfem::IntegrationRule& rule = ::mfem::MassIntegrator::GetRule(element, element, transformation);
            for (int point_index = 0; point_index < rule.GetNPoints(); ++point_index) {
                const ::mfem::IntegrationPoint& point = rule.IntPoint(point_index);
                transformation.SetIntPoint(&point);
                element.CalcShape(point, shape);
                long double norm_squared = 0.0L;
                for (int component = 0; component < 3; ++component) {
                    long double value = 0.0L;
                    for (int node = 0; node < element.GetDof(); ++node) {
                        value += static_cast<long double>(shape[node]) * local[node + component * element.GetDof()];
                    }
                    norm_squared += value * value;
                }
                work += point.weight * transformation.Weight() * norm_squared;
            }
        }
        REQUIRE(std::isfinite(work));
        REQUIRE(work > 0.0L);
        return static_cast<double>(work);
    }
}

TEST_CASE("Harmonic reference metrics measure the actual vector finite-element surface trace", "[integration][deformation]") {
    for (double flattening : {0.0, 0.2}) {
        // MeanField uses cubic displacement traces; refine the surface to resolve all 25 degree-four modes.
        PairedMesh fixture(1, flattening, 3);
        const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
        REQUIRE(reference);
        for (int cutoff : {0, 4}) {
            const auto embedding = deformation::build_surface_embedding(deformation::SphericalHarmonicSurface{.maximum_degree = 4, .maximum_azimuthal_order = cutoff}, reference->positions);
            REQUIRE(embedding);
            const auto metric = backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.surface_attribute, *reference, *embedding);
            REQUIRE(metric);
            REQUIRE(metric->size() == static_cast<std::size_t>(cutoff == 0 ? 5 : 25));
            CHECK((*metric)[0] == Catch::Approx(analytic_trace_norm(fixture, 0)).epsilon(5e-13));
            CHECK((*metric)[cutoff == 0 ? 1 : 2] == Catch::Approx(analytic_trace_norm(fixture, 1)).epsilon(5e-13));
            for (double entry : *metric) {
                CHECK(std::isfinite(entry));
                CHECK(entry > 0.0);
            }
        }
    }
}

TEST_CASE("Nodal radial metrics match the independent assembled boundary mass diagonal", "[integration][deformation]") {
    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    const auto embedding = deformation::build_surface_embedding(deformation::NodalRadialSurface{}, reference->positions);
    REQUIRE(embedding);
    const auto metric = backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.surface_attribute, *reference, *embedding);
    REQUIRE(metric);
    ::mfem::Array<int> marker(fixture.generated.mesh->bdr_attributes.Max());
    marker = 0;
    marker[fixture.surface_attribute - 1] = 1;
    ::mfem::BilinearForm mass(&fixture.scalar_space);
    mass.AddBoundaryIntegrator(new ::mfem::MassIntegrator(), marker);
    mass.Assemble();
    mass.Finalize();
    ::mfem::Vector diagonal;
    mass.SpMat().GetDiag(diagonal);
    REQUIRE(metric->size() == reference->scalar_dofs.size());
    for (std::size_t node = 0; node < metric->size(); ++node) {
        CHECK((*metric)[node] == Catch::Approx(diagonal[reference->scalar_dofs[node]]).epsilon(5e-14));
    }
}

TEST_CASE("Reference metric compilation rejects incompatible trace ordering and geometry", "[integration][deformation]") {
    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    const auto embedding = deformation::build_surface_embedding(deformation::SphericalHarmonicSurface{.maximum_degree = 2}, reference->positions);
    REQUIRE(embedding);
    deformation::ReferenceSurface changed = *reference;
    SECTION("Canonical scalar node ordering changed") {
        REQUIRE(changed.scalar_dofs.size() > 1);
        std::swap(changed.scalar_dofs[0], changed.scalar_dofs[1]);
    }
    SECTION("Reference positions changed at fixed dimensions") {
        changed.positions[0][0] += 0.01;
    }
    SECTION("A reference node is missing") {
        changed.positions.pop_back();
        changed.scalar_dofs.pop_back();
    }
    const auto result = backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.surface_attribute, changed, *embedding);
    REQUIRE_FALSE(result);
    CHECK(result.error().reason == deformation::DeformationFailureReason::incompatible_discretization);
}

TEST_CASE("Reference metrics reject an embedding from different points despite equal dimensions", "[integration][deformation]") {
    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    auto positions = reference->positions;
    positions[0][0] += 0.01;
    const auto embedding = deformation::build_surface_embedding(deformation::NodalRadialSurface{}, positions);
    REQUIRE(embedding);
    const auto result = backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.surface_attribute, *reference, *embedding);
    REQUIRE_FALSE(result);
    CHECK(result.error().reason == deformation::DeformationFailureReason::incompatible_discretization);
}

TEST_CASE("Reference metrics reject a changed mesh before assembling parameter norms", "[integration][deformation]") {
    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    const auto embedding = deformation::build_surface_embedding(deformation::NodalRadialSurface{}, reference->positions);
    REQUIRE(embedding);
    /// Moving the reference coordinates must invalidate the old trace even if all vector dimensions stay equal.
    ::mfem::VectorFunctionCoefficient translation(3, [](const ::mfem::Vector& old_position, ::mfem::Vector& new_position) {
        new_position = old_position;
        new_position[0] += 0.01;
    });
    fixture.generated.mesh->Transform(translation);
    const auto result = backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.surface_attribute, *reference, *embedding);
    REQUIRE_FALSE(result);
    CHECK(result.error().reason == deformation::DeformationFailureReason::incompatible_discretization);
}

TEST_CASE("Reference metrics require the selected scalar surface trace", "[integration][deformation]") {
    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    const auto embedding = deformation::build_surface_embedding(deformation::NodalRadialSurface{}, reference->positions);
    REQUIRE(embedding);
    CHECK_FALSE(backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.infinity_attribute, *reference, *embedding));
    CHECK_FALSE(backend::surface_parameter_mass_diagonal(fixture.scalar_space, 999, *reference, *embedding));
    CHECK_FALSE(backend::surface_parameter_mass_diagonal(fixture.vector_space, fixture.surface_attribute, *reference, *embedding));
    const deformation::ReferenceSurface copied = *reference;
    CHECK(backend::surface_parameter_mass_diagonal(fixture.scalar_space, fixture.surface_attribute, copied, *embedding));
}

TEST_CASE("Harmonic and nodal surfaces couple to identical stellar and vacuum volume fields", "[integration][deformation]") {
    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    auto harmonic = deformation::build_surface_embedding(deformation::SphericalHarmonicSurface{.maximum_degree = 2}, reference->positions);
    auto nodal = deformation::build_surface_embedding(deformation::NodalRadialSurface{}, reference->positions);
    REQUIRE(harmonic);
    REQUIRE(nodal);
    auto interpolation = fixture.interpolation();
    REQUIRE(interpolation);
    auto extension = deformation::make_radial_extension(std::make_shared<const deformation::RadialInterpolation>(std::move(*interpolation)));
    REQUIRE(extension);
    const auto harmonic_map = std::make_shared<const deformation::SurfaceEmbedding>(std::move(*harmonic));
    const auto nodal_map = std::make_shared<const deformation::SurfaceEmbedding>(std::move(*nodal));
    const auto volume_extension = std::make_shared<const deformation::RadialExtension>(std::move(*extension));
    const auto harmonic_domain = deformation::compose_deformation(harmonic_map, volume_extension);
    const auto nodal_domain = deformation::compose_deformation(nodal_map, volume_extension);
    REQUIRE(harmonic_domain);
    REQUIRE(nodal_domain);

    std::vector<double> harmonic_parameters(harmonic_map->input_size());
    for (std::size_t parameter = 0; parameter < harmonic_parameters.size(); ++parameter) {
        harmonic_parameters[parameter] = 0.001 * std::sin(0.37 * static_cast<double>(parameter + 1));
    }
    std::vector<double> nodal_parameters(nodal_map->input_size());
    for (std::size_t node = 0; node < nodal_parameters.size(); ++node) {
        long double value = 0.0L;
        for (std::size_t parameter = 0; parameter < harmonic_parameters.size(); ++parameter) {
            value += static_cast<long double>(harmonic_map->nodal_basis()[node * harmonic_parameters.size() + parameter]) * harmonic_parameters[parameter];
        }
        nodal_parameters[node] = static_cast<double>(value);
    }
    const auto harmonic_volume = deformation::evaluate_displacement(*harmonic_domain, harmonic_parameters);
    const auto nodal_volume = deformation::evaluate_displacement(*nodal_domain, nodal_parameters);
    REQUIRE(harmonic_volume);
    REQUIRE(nodal_volume);
    REQUIRE(harmonic_volume->size() == nodal_volume->size());
    long double difference_norm = 0.0L;
    long double reference_norm = 0.0L;
    for (std::size_t entry = 0; entry < harmonic_volume->size(); ++entry) {
        const long double difference = (*harmonic_volume)[entry] - (*nodal_volume)[entry];
        difference_norm += difference * difference;
        reference_norm += static_cast<long double>((*nodal_volume)[entry]) * (*nodal_volume)[entry];
    }
    REQUIRE(reference_norm > 0.0L);
    CHECK(std::sqrt(difference_norm / reference_norm) < 3e-13);
    const auto infinity_nodes = backend::select_surface_scalar_dofs(fixture.scalar_space, fixture.infinity_attribute);
    REQUIRE(infinity_nodes);
    for (int node : *infinity_nodes) {
        for (int component = 0; component < 3; ++component) {
            CHECK(std::abs((*harmonic_volume)[node + component * fixture.scalar_space.GetNDofs()]) < 1e-15);
        }
    }
}

TEST_CASE("Shared boundary traversal rejects malformed coordinates before numerical assembly", "[integration][deformation][backend]") {
    namespace shared = serif::backend::mfem::shared;

    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    std::vector<int> selected = reference->scalar_dofs;
    REQUIRE_FALSE(selected.empty());
    SECTION("Repeated scalar nodes") {
        selected.push_back(selected.front());
    }
    SECTION("An out-of-range scalar node") {
        selected.front() = fixture.scalar_space.GetNDofs();
    }
    SECTION("A negative scalar node") {
        selected.front() = -1;
    }
    SECTION("No selected coordinates for an existing boundary") {
        selected.clear();
    }
    bool assembled = false;
    const shared::BoundaryStatus status = shared::for_each_boundary_element(fixture.scalar_space, fixture.surface_attribute, selected,
        [&](const shared::BoundaryElement&) -> shared::BoundaryStatus {
            assembled = true;
            return {};
        });
    REQUIRE_FALSE(status);
    CHECK(status.error().incompatible_coordinates);
    CHECK_FALSE(assembled);
}

TEST_CASE("Shared boundary traversal stops and identifies a rejected element assembly", "[integration][deformation][backend]") {
    namespace shared = serif::backend::mfem::shared;

    PairedMesh fixture;
    const auto reference = backend::sample_reference_surface(fixture.scalar_space, fixture.surface_attribute);
    REQUIRE(reference);
    int visited = 0;
    const shared::BoundaryStatus status = shared::for_each_boundary_element(fixture.scalar_space, fixture.surface_attribute, reference->scalar_dofs,
        [&](const shared::BoundaryElement&) -> shared::BoundaryStatus {
            ++visited;
            return std::unexpected(shared::BoundaryFailure{.message = "Rejected local numerical contribution.", .incompatible_coordinates = false});
        });
    REQUIRE_FALSE(status);
    CHECK(visited == 1);
    CHECK_FALSE(status.error().incompatible_coordinates);
    CHECK(status.error().message == "Rejected local numerical contribution.");
    REQUIRE(status.error().element >= 0);
    CHECK(fixture.scalar_space.GetMesh()->GetBdrAttribute(status.error().element) == fixture.surface_attribute);
}