#include <algorithm>
#include <cmath>
#include <limits>
#include <numeric>
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_approx.hpp>
#include "mfem/fem/gridfunc.hpp"
#include "serif/tests/deformation/mesh.hpp"

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

    /** @brief Expose a numbering mismatch without requiring unsupported variable-order nonconforming meshes. */
    struct ReorderedCompanionSpace : ::mfem::FiniteElementSpace
    {
        using ::mfem::FiniteElementSpace::FiniteElementSpace;

        /** @brief Preserve the basis, collection and total size while swapping two local scalar coefficients. */
        void GetElementDofs(int element, ::mfem::Array<int>& dofs, ::mfem::DofTransformation& transformation) const override
        {
            ::mfem::FiniteElementSpace::GetElementDofs(element, dofs, transformation);
            std::swap(dofs[0], dofs[1]);
        }
    };

    /** @brief Keep action tests concise while ensuring compilation failure is reported at its source. */
    [[nodiscard]] std::shared_ptr<const deformation::RadialInterpolation> interpolation(PairedMesh& mesh)
    {
        auto result = mesh.interpolation();
        if (!result)
        {
            INFO(result.error().message);
            REQUIRE(result.has_value());
        }
        return std::make_shared<const deformation::RadialInterpolation>(std::move(*result));
    }

    /** @brief Preserve the MeanField finite-difference normalization, including its epsilon floor. */
    [[nodiscard]] double relative_error(std::span<const double> actual, std::span<const double> expected)
    {
        double squared_error = 0.0, squared_norm = 0.0;
        for (std::size_t entry = 0; entry < actual.size(); ++entry)
        {
            const double difference = actual[entry] - expected[entry];
            squared_error += difference * difference;
            squared_norm += expected[entry] * expected[entry];
        }
        return std::sqrt(squared_error) / std::max(std::sqrt(squared_norm), std::numeric_limits<double>::epsilon());
    }
} // namespace

TEST_CASE("Radial interpolation preserves constants on paired logical rays and canonical physical surface nodes", "[unit][deformation]")
{
    // Cubic nodes on a refined core project between surface nodes; the unrefined quadratic fixture is exactly aligned.
    PairedMesh mesh(1, 0.0, 3);
    const auto data = interpolation(mesh);
    const auto logical_positions = backend::sample_reference_positions(mesh.scalar_space, *mesh.generated.reference_mesh);
    REQUIRE(logical_positions.has_value());
    const auto surface = backend::sample_reference_surface(mesh.scalar_space, PairedMesh::surface_attribute);
    REQUIRE(surface.has_value());
    CHECK(surface->scalar_dofs == data->data().surface_scalar_dofs);
    CHECK(std::is_sorted(surface->scalar_dofs.begin(), surface->scalar_dofs.end()));
    CHECK(surface->positions.size() == data->surface_node_count());
    CHECK(data->data().stellar_radius > 0.0);
    CHECK(data->data().infinity_radius > data->data().stellar_radius);
    CHECK(data->scalar_dof_count() == static_cast<std::size_t>(mesh.scalar_space.GetNDofs()));

    int interpolated_rows = 0;
    int centers = 0;
    for (std::size_t node = 0; node < data->scalar_dof_count(); ++node)
    {
        const int begin = data->data().row_offsets[node];
        const int end = data->data().row_offsets[node + 1];
        if (data->data().logical_radii[node] < 1.0e-14)
        {
            CHECK(begin == end);
            ++centers;
            continue;
        }
        const double sum = std::accumulate(data->data().interpolation_weights.begin() + begin, data->data().interpolation_weights.begin() + end, 0.0);
        CHECK(sum == Catch::Approx(1.0).margin(2.0e-12));
        for (std::size_t component = 0; component < 3; ++component)
        {
            double interpolated_position = 0.0;
            for (int entry = begin; entry < end; ++entry)
            {
                const int surface_dof = data->data().surface_scalar_dofs[data->data().surface_columns[entry]];
                interpolated_position += data->data().interpolation_weights[entry] * (*logical_positions)[surface_dof][component];
            }
            const double projected_position = data->data().stellar_radius * (*logical_positions)[node][component] / data->data().logical_radii[node];
            CHECK(interpolated_position == Catch::Approx(projected_position).margin(2.0e-11));
        }
        if (end - begin > 1)
        {
            ++interpolated_rows;
        }
    }
    CHECK(centers > 0);
    CHECK(interpolated_rows > 0);
}

TEST_CASE("Fused radial extension reproduces the surface fixes center and infinity and preserves original powers", "[unit][deformation]")
{
    PairedMesh mesh;
    const auto data = interpolation(mesh);
    const auto quadratic = deformation::make_radial_extension(data);
    const auto cubic = deformation::make_radial_extension(data, {.radial_power = 3.0});
    REQUIRE(quadratic.has_value());
    REQUIRE(cubic.has_value());
    CHECK(quadratic->interpolation() == cubic->interpolation());
    std::vector<double> surface(quadratic->input_size());
    for (std::size_t node = 0; node < data->surface_node_count(); ++node)
    {
        surface[3 * node] = 0.02;
        surface[3 * node + 1] = -0.03;
        surface[3 * node + 2] = 0.05;
    }
    std::vector<double> output(quadratic->output_size());
    REQUIRE(quadratic->apply(surface, output));
    int shared_nodes = 0, fixed_centers = 0, fixed_infinity = 0;
    for (std::size_t node = 0; node < data->scalar_dof_count(); ++node)
    {
        const double radius = data->data().logical_radii[node];
        double weight = 0.0;
        if (data->data().stellar_support[node])
        {
            weight = std::pow(radius / data->data().stellar_radius, 2.0);
            CHECK(cubic->radial_weights()[node] == Catch::Approx(std::pow(radius / data->data().stellar_radius, 3.0)).margin(2.0e-14));
        }
        else if (data->data().vacuum_support[node])
        {
            weight = (data->data().infinity_radius - radius) / (data->data().infinity_radius - data->data().stellar_radius);
            CHECK(cubic->radial_weights()[node] == quadratic->radial_weights()[node]);
        }
        for (std::size_t component = 0; component < 3; ++component)
        {
            CHECK(output[node + component * data->scalar_dof_count()] == Catch::Approx(weight * surface[component]).margin(2.0e-13));
        }
        if (radius < 1.0e-14)
        {
            CHECK(quadratic->radial_weights()[node] == 0.0);
            ++fixed_centers;
        }
        if (std::abs(radius - data->data().infinity_radius) < 1.0e-14)
        {
            CHECK(quadratic->radial_weights()[node] == Catch::Approx(0.0).margin(1.0e-14));
            ++fixed_infinity;
        }
        if (data->data().stellar_support[node] && data->data().vacuum_support[node])
        {
            CHECK(quadratic->radial_weights()[node] == Catch::Approx(1.0).margin(2.0e-14));
            ++shared_nodes;
        }
    }
    CHECK(fixed_centers > 0);
    CHECK(fixed_infinity > 0);
    CHECK(shared_nodes == static_cast<int>(data->surface_node_count()));

    for (std::size_t entry = 0; entry < surface.size(); ++entry)
    {
        surface[entry] = std::sin(0.17 * (entry + 1));
    }
    REQUIRE(quadratic->apply(surface, output));
    for (std::size_t node = 0; node < data->surface_node_count(); ++node)
    {
        for (std::size_t component = 0; component < 3; ++component)
        {
            CHECK(output[data->data().surface_scalar_dofs[node] + component * data->scalar_dof_count()] == Catch::Approx(surface[3 * node + component]).margin(2.0e-12));
        }
    }
}

TEST_CASE("Radial extension derivatives retain centered differences and exact transpose work after refinement", "[unit][deformation]")
{
    for (const int refinement : {0, 1})
    {
        CAPTURE(refinement);
        PairedMesh mesh(refinement, 0.1);
        const auto extension = deformation::make_radial_extension(interpolation(mesh));
        REQUIRE(extension.has_value());
        std::vector<double> direction(extension->input_size()), plus(direction.size()), minus(direction.size());
        constexpr double step = 1.0e-6;
        for (std::size_t entry = 0; entry < direction.size(); ++entry)
        {
            const double index = entry + 1;
            direction[entry] = std::cos(0.13 * index) - 0.2 * std::sin(0.31 * index);
            const double surface = 0.01 * std::sin(0.17 * index);
            plus[entry] = surface + step * direction[entry];
            minus[entry] = surface - step * direction[entry];
        }
        std::vector<double> plus_output(extension->output_size()), minus_output(plus_output.size()), derivative(plus_output.size()), force(plus_output.size()), surface_force(direction.size());
        REQUIRE(extension->apply(plus, plus_output));
        REQUIRE(extension->apply(minus, minus_output));
        REQUIRE(extension->apply(direction, derivative));
        for (std::size_t entry = 0; entry < derivative.size(); ++entry)
        {
            plus_output[entry] = (plus_output[entry] - minus_output[entry]) / (2.0 * step);
            const double index = entry + 1;
            force[entry] = std::sin(0.07 * index) + 0.3 * std::cos(0.11 * index);
        }
        CHECK(relative_error(derivative, plus_output) < 2.0e-10);
        REQUIRE(extension->apply_transpose(force, surface_force));
        const double volume_work = std::inner_product(derivative.begin(), derivative.end(), force.begin(), 0.0);
        const double surface_work = std::inner_product(direction.begin(), direction.end(), surface_force.begin(), 0.0);
        CHECK(std::abs(volume_work - surface_work) <= 5.0e-13 * std::max({1.0, std::abs(volume_work), std::abs(surface_work)}));
    }
}

TEST_CASE("Logical radial interpolation is independent of physical flattening", "[unit][deformation]")
{
    PairedMesh spherical;
    PairedMesh oblate(0, 0.15);
    const auto spherical_data = interpolation(spherical);
    const auto oblate_data = interpolation(oblate);
    CHECK(spherical_data->data().row_offsets == oblate_data->data().row_offsets);
    CHECK(spherical_data->data().surface_columns == oblate_data->data().surface_columns);
    CHECK(spherical_data->data().interpolation_weights == oblate_data->data().interpolation_weights);
    CHECK(spherical_data->data().logical_radii == oblate_data->data().logical_radii);
    const auto spherical_surface = backend::sample_reference_surface(spherical.scalar_space, PairedMesh::surface_attribute);
    const auto oblate_surface = backend::sample_reference_surface(oblate.scalar_space, PairedMesh::surface_attribute);
    REQUIRE(spherical_surface.has_value());
    REQUIRE(oblate_surface.has_value());
    CHECK(spherical_surface->positions != oblate_surface->positions);
}

TEST_CASE("Radial compilation rejects unsupported spaces broken mesh pairing and invalid foliation", "[unit][deformation]")
{
    PairedMesh mesh;
    SECTION("Wrong vector storage")
    {
        ::mfem::FiniteElementSpace interleaved(mesh.generated.mesh.get(), &mesh.collection, 3, ::mfem::Ordering::byVDIM);
        CHECK_FALSE(backend::build_radial_interpolation(mesh.scalar_space, interleaved, *mesh.generated.reference_mesh, mesh.stellar_elements, mesh.vacuum_elements, 1, 2));
    }
    SECTION("Different companion collection")
    {
        ::mfem::H1_FECollection different_collection(2, 3);
        ::mfem::FiniteElementSpace different_vector(mesh.generated.mesh.get(), &different_collection, 3);
        CHECK_FALSE(backend::build_radial_interpolation(mesh.scalar_space, different_vector, *mesh.generated.reference_mesh, mesh.stellar_elements, mesh.vacuum_elements, 1, 2));
    }
    SECTION("Matching sizes and collection with different scalar coefficient numbering")
    {
        ReorderedCompanionSpace reordered(mesh.generated.mesh.get(), &mesh.collection, 3, ::mfem::Ordering::byNODES);
        REQUIRE(reordered.GetVSize() == mesh.vector_space.GetVSize());
        REQUIRE(reordered.FEColl() == mesh.scalar_space.FEColl());
        const auto result = backend::build_radial_interpolation(mesh.scalar_space, reordered, *mesh.generated.reference_mesh, mesh.stellar_elements, mesh.vacuum_elements, 1, 2);
        REQUIRE_FALSE(result);
        CHECK(result.error().reason == deformation::DeformationFailureReason::incompatible_discretization);
        CHECK(result.error().element == 0);
    }
    SECTION("Discontinuous scalar basis")
    {
        ::mfem::L2_FECollection discontinuous(2, 3);
        ::mfem::FiniteElementSpace scalar(mesh.generated.mesh.get(), &discontinuous);
        CHECK_FALSE(backend::sample_reference_surface(scalar, 1));
    }
    SECTION("Missing boundary")
    {
        CHECK_FALSE(backend::sample_reference_surface(mesh.scalar_space, 77));
        CHECK_FALSE(backend::build_radial_interpolation(mesh.scalar_space, mesh.vector_space, *mesh.generated.reference_mesh, mesh.stellar_elements, mesh.vacuum_elements, 1, 77));
    }
    SECTION("Malformed element support")
    {
        mesh.stellar_elements.pop_back();
        CHECK_FALSE(mesh.interpolation());
    }
    SECTION("Non-binary element support")
    {
        mesh.stellar_elements[0] = -1;
        CHECK_FALSE(mesh.interpolation());
    }
    SECTION("Reversed logical shells")
    {
        CHECK_FALSE(backend::build_radial_interpolation(mesh.scalar_space, mesh.vector_space, *mesh.generated.reference_mesh, mesh.stellar_elements, mesh.vacuum_elements, 2, 1));
    }
    SECTION("Inconsistent paired boundary attributes")
    {
        mesh.generated.reference_mesh->GetBdrElement(0)->SetAttribute(77);
        CHECK_FALSE(mesh.interpolation());
    }
    SECTION("Inconsistent paired element attributes")
    {
        mesh.generated.reference_mesh->GetElement(0)->SetAttribute(77);
        CHECK_FALSE(mesh.interpolation());
    }
    SECTION("Inconsistent paired element orientation")
    {
        int* vertices = mesh.generated.reference_mesh->GetElement(0)->GetVertices();
        std::swap(vertices[0], vertices[1]);
        CHECK_FALSE(mesh.interpolation());
    }
    SECTION("Physical geometry is not a logical cubical foliation")
    {
        CHECK_FALSE(backend::build_radial_interpolation(mesh.scalar_space, mesh.vector_space, *mesh.generated.mesh, mesh.stellar_elements, mesh.vacuum_elements, 1, 2));
    }
    SECTION("Vacuum support incorrectly includes the center")
    {
        std::fill(mesh.vacuum_elements.begin(), mesh.vacuum_elements.end(), 1);
        CHECK_FALSE(mesh.interpolation());
    }
}

TEST_CASE("Radial actions reject invalid values lengths and overlapping buffers", "[unit][deformation]")
{
    PairedMesh mesh;
    const auto data = interpolation(mesh);
    const auto extension = deformation::make_radial_extension(data);
    REQUIRE(extension.has_value());
    CHECK_FALSE(deformation::make_radial_extension(data, {.radial_power = 0.5}));
    CHECK_FALSE(deformation::make_radial_extension(data, {.radial_power = std::numeric_limits<double>::infinity()}));
    CHECK_FALSE(deformation::make_radial_extension(nullptr));
    std::vector<double> surface(extension->input_size(), 1.0), volume(extension->output_size(), 7.0);
    CHECK_FALSE(extension->apply(std::span(surface).first(surface.size() - 1), volume));
    CHECK(std::all_of(volume.begin(), volume.end(), [](double value)
    {
        return value == 7.0;
    }));
    surface[0] = std::numeric_limits<double>::quiet_NaN();
    CHECK_FALSE(extension->apply(surface, volume));
    surface[0] = 1.0;
    volume[0] = std::numeric_limits<double>::infinity();
    CHECK_FALSE(extension->apply_transpose(volume, surface));
    std::vector<double> overlap(extension->input_size() + extension->output_size(), 1.0);
    CHECK_FALSE(extension->apply(std::span<const double>(overlap).first(extension->input_size()), std::span<double>(overlap).first(extension->output_size())));
    CHECK_FALSE(extension->apply_transpose(std::span<const double>(overlap).first(extension->output_size()), std::span<double>(overlap).first(extension->input_size())));
    volume[0] = 7.0;
    REQUIRE(extension->apply(surface, volume));
    REQUIRE(extension->apply_transpose(volume, surface));
}

TEST_CASE("Imported radial interpolation is checked once and owns its validated data", "[unit][deformation]")
{
    PairedMesh mesh;
    const auto data = interpolation(mesh);
    deformation::RadialInterpolationData malformed = data->data();
    malformed.surface_columns.front() = static_cast<int>(data->surface_node_count());
    CHECK_FALSE(deformation::make_radial_interpolation(malformed));
    malformed = data->data();
    malformed.row_offsets.back() += 1;
    CHECK_FALSE(deformation::make_radial_interpolation(malformed));
    malformed = data->data();
    const int surface_dof = data->data().surface_scalar_dofs.front();
    for (int entry = malformed.row_offsets[surface_dof]; entry < malformed.row_offsets[surface_dof + 1]; ++entry)
    {
        malformed.interpolation_weights[entry] *= 0.5;
    }
    CHECK_FALSE(deformation::make_radial_interpolation(malformed));

    malformed = data->data();
    const int first_entry = malformed.row_offsets[surface_dof];
    const int next_entry = malformed.row_offsets[surface_dof + 1];
    const int removed_entries = next_entry - first_entry;
    REQUIRE(removed_entries > 0);
    malformed.surface_columns.erase(malformed.surface_columns.begin() + first_entry, malformed.surface_columns.begin() + next_entry);
    malformed.interpolation_weights.erase(malformed.interpolation_weights.begin() + first_entry, malformed.interpolation_weights.begin() + next_entry);
    for (std::size_t row = surface_dof + 1; row < malformed.row_offsets.size(); ++row)
    {
        malformed.row_offsets[row] -= removed_entries;
    }
    CHECK_FALSE(deformation::make_radial_interpolation(malformed));
    SECTION("Nonfinite interpolation weight")
    {
        malformed = data->data();
        malformed.interpolation_weights.front() = std::numeric_limits<double>::quiet_NaN();
        CHECK_FALSE(deformation::make_radial_interpolation(malformed));
    }
    SECTION("Invalid material support")
    {
        malformed = data->data();
        malformed.stellar_support.front() = 2;
        CHECK_FALSE(deformation::make_radial_interpolation(malformed));
    }
    SECTION("Nonfinite logical radius")
    {
        malformed = data->data();
        malformed.logical_radii.front() = std::numeric_limits<double>::infinity();
        CHECK_FALSE(deformation::make_radial_interpolation(malformed));
    }
    SECTION("Imported storage cannot mutate a compiled table")
    {
        deformation::RadialInterpolationData imported = data->data();
        const auto checked = deformation::make_radial_interpolation(imported);
        REQUIRE(checked);
        imported.interpolation_weights.front() = std::numeric_limits<double>::quiet_NaN();
        CHECK(checked->data().interpolation_weights == data->data().interpolation_weights);
        const auto shared = std::make_shared<const deformation::RadialInterpolation>(*checked);
        const auto first = deformation::make_radial_extension(shared);
        const auto second = deformation::make_radial_extension(shared, {.radial_power = 3.0});
        REQUIRE(first);
        REQUIRE(second);
        CHECK(first->interpolation() == shared);
        CHECK(second->interpolation() == shared);
        CHECK_FALSE(std::ranges::equal(first->radial_weights(), second->radial_weights()));
    }
}

TEST_CASE("Extended constant vectors retain scalar finite element interpolation between volume nodes", "[unit][deformation]")
{
    PairedMesh mesh;
    const auto extension = deformation::make_radial_extension(interpolation(mesh));
    REQUIRE(extension.has_value());
    constexpr std::array<double, 3> components{1.25, -0.75, 2.5};
    std::vector<double> surface(extension->input_size()), volume(extension->output_size());
    for (std::size_t node = 0; node < surface.size() / 3; ++node)
    {
        for (std::size_t component = 0; component < 3; ++component)
        {
            surface[3 * node + component] = components[component];
        }
    }
    REQUIRE(extension->apply(surface, volume));
    ::mfem::GridFunction displacement(&mesh.vector_space);
    ::mfem::GridFunction radial_weight(&mesh.scalar_space);
    std::copy(volume.begin(), volume.end(), displacement.GetData());
    std::copy(extension->radial_weights().begin(), extension->radial_weights().end(), radial_weight.GetData());
    ::mfem::Vector value(3);
    for (int element = 0; element < mesh.generated.mesh->GetNE(); ++element)
    {
        const ::mfem::IntegrationPoint& point = ::mfem::Geometries.GetCenter(mesh.generated.mesh->GetElementGeometry(element));
        const double weight = radial_weight.GetValue(element, point);
        displacement.GetVectorValue(element, point, value);
        for (int component = 0; component < 3; ++component)
        {
            CHECK(value[component] == Catch::Approx(components[component] * weight).margin(2.0e-11));
        }
    }
}

TEST_CASE("Small radial surface motion preserves orientation before and after uniform refinement", "[unit][deformation]")
{
    for (const int refinement : {0, 1})
    {
        CAPTURE(refinement);
        PairedMesh mesh(refinement);
        const auto surface_geometry = backend::sample_reference_surface(mesh.scalar_space, PairedMesh::surface_attribute);
        REQUIRE(surface_geometry.has_value());
        const auto extension = deformation::make_radial_extension(interpolation(mesh));
        REQUIRE(extension.has_value());
        std::vector<double> surface(extension->input_size()), volume(extension->output_size());
        for (std::size_t node = 0; node < surface_geometry->positions.size(); ++node)
        {
            const std::array<double, 3>& position = surface_geometry->positions[node];
            const double radius = std::sqrt(position[0] * position[0] + position[1] * position[1] + position[2] * position[2]);
            const double polar_direction = position[2] / radius;
            const double amplitude = 0.006 - 0.001 * 0.5 * (3.0 * polar_direction * polar_direction - 1.0);
            for (std::size_t component = 0; component < 3; ++component)
            {
                surface[3 * node + component] = amplitude * position[component] / radius;
            }
        }
        REQUIRE(extension->apply(surface, volume));
        ::mfem::GridFunction displacement(&mesh.vector_space);
        std::copy(volume.begin(), volume.end(), displacement.GetData());
        double minimum_determinant = std::numeric_limits<double>::infinity();
        ::mfem::DenseMatrix gradient;
        for (int element = 0; element < mesh.generated.mesh->GetNE(); ++element)
        {
            ::mfem::ElementTransformation& transformation = *mesh.generated.mesh->GetElementTransformation(element);
            const ::mfem::IntegrationRule& rule = ::mfem::IntRules.Get(transformation.GetGeometryType(), mesh.vector_space.GetFE(element)->GetOrder() + 2);
            for (int point = 0; point < rule.GetNPoints(); ++point)
            {
                transformation.SetIntPoint(&rule.IntPoint(point));
                displacement.GetVectorGradient(transformation, gradient);
                for (int component = 0; component < 3; ++component)
                {
                    gradient(component, component) += 1.0;
                }
                minimum_determinant = std::min(minimum_determinant, gradient.Det());
            }
        }
        CHECK(minimum_determinant > 0.0);
    }
}