// MeanField field_dof_map.cpp and field_mfem.cpp invariants expressed through serial TermBinding.
#include <algorithm>
#include <array>
#include <limits>
#include <memory>
#include <vector>

#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/gridfunc.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/backend/mfem/discretization/dof_selection.hpp"
#include "serif/discretization/domain/schema/schemas.hpp"
#include "serif/tests/deformation/mesh.hpp"

namespace
{
    namespace discretization = serif::discretization;
    namespace backend = discretization::backend::mfem;
    namespace blocks = discretization::blocks;
    namespace domain = discretization::domain;
    using Schema = domain::schema::CoreEnvelopeVacuumDomainSchema;
    using AlternateSchema = domain::schema::DomainSchema<domain::ids::lists::DomainIDList<domain::ids::DomainID<domain::CoreDomain, 11>, domain::ids::DomainID<domain::EnvelopeDomain, 17>, domain::ids::DomainID<domain::VacuumDomain, 29>>, domain::ids::lists::BoundaryIDList<>, domain::relation::lists::RelationList<>>;

    /** @brief A stellar velocity represented by independent discontinuous scalar components. */
    struct DiscontinuousVelocity final : blocks::field
    {
        struct vector final : blocks::term<DiscontinuousVelocity>
        {
            using FunctionalSpace = discretization::functional_space::DiscreteFunctionalSpace<discretization::functional_space::L2, 2, 1>;
            using Support = domain::StellarDomains;
        };
    };

    /** @brief Two adjacent cells expose both inactive vacuum DOFs and shared continuous trace DOFs. */
    [[nodiscard]] mfem::Mesh split_mesh(const int stellar_attribute = 2, const int vacuum_attribute = 3)
    {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON, 2.0, 1.0, 1.0);
        mesh.GetElement(0)->SetAttribute(stellar_attribute);
        mesh.GetElement(1)->SetAttribute(vacuum_attribute);
        mesh.SetAttributes();
        return mesh;
    }

    /** @brief Independently mark the native element DOFs selected by semantic support. */
    template <typename Term, typename DomainSchema>
    [[nodiscard]] std::vector<bool> expected_support(const mfem::FiniteElementSpace& space)
    {
        std::vector<bool> active(space.GetVSize(), false);
        mfem::Array<int> dofs;
        for (int element = 0; element < space.GetNE(); ++element)
        {
            if (!DomainSchema::template domain_id_belongs_to<typename Term::Support>(space.GetMesh()->GetAttribute(element)))
            {
                continue;
            }
            space.GetElementVDofs(element, dofs);
            for (const int encoded : dofs)
            {
                active[encoded < 0 ? -1 - encoded : encoded] = true;
            }
        }
        return active;
    }

    /** @brief Compare canonical indexing, projection, and GridFunction transfer with the independent marker. */
    template <typename Term, typename DomainSchema = Schema>
    void check_supported_projection(mfem::FiniteElementSpace& space)
    {
        const backend::TermBinding<Term, DomainSchema> binding{space};
        const std::vector<bool> active = expected_support<Term, DomainSchema>(space);
        REQUIRE(binding.full_size() == space.GetVSize());
        CHECK(binding.size() == std::ranges::count(active, true));
        mfem::Vector full(binding.full_size());
        for (int dof = 0; dof < full.Size(); ++dof)
        {
            full[dof] = 10.0 + dof;
        }
        const mfem::Vector reduced = binding.gather(full);
        const mfem::Vector projected = binding.scatter(reduced);
        const mfem::Vector recovered = binding.gather(projected);
        int next_reduced = 0;
        for (int dof = 0; dof < full.Size(); ++dof)
        {
            CAPTURE(dof);
            const std::optional<int> index = binding.reduced_dof(dof);
            CHECK(index.has_value() == active[dof]);
            if (active[dof])
            {
                REQUIRE(index);
                CHECK(*index == next_reduced++);
                CHECK(reduced[*index] == full[dof]);
                CHECK(recovered[*index] == reduced[*index]);
                CHECK(projected[dof] == full[dof]);
            }
            else
            {
                CHECK(projected[dof] == 0.0);
            }
        }

        // MFEM owns the GridFunction conversion; no additional field-adapter lifecycle is needed.
        mfem::GridFunction grid(&space);
        grid.SetFromTrueDofs(projected);
        mfem::Vector true_values;
        grid.GetTrueDofs(true_values);
        const mfem::Vector grid_values = binding.gather(true_values);
        for (int dof = 0; dof < reduced.Size(); ++dof)
        {
            CHECK(grid_values[dof] == reduced[dof]);
        }
        std::vector<double> storage(full.Size() + 2, -77.0);
        std::copy_n(full.GetData(), full.Size(), storage.data() + 1);
        const mfem::Vector full_view(storage.data() + 1, full.Size());
        const mfem::Vector from_view = binding.gather(full_view);
        for (int dof = 0; dof < reduced.Size(); ++dof)
        {
            CHECK(from_view[dof] == reduced[dof]);
        }
        CHECK(storage.front() == -77.0);
        CHECK(storage.back() == -77.0);
        CHECK_THROWS_AS(binding.reduced_dof(-1), std::out_of_range);
        CHECK_THROWS_AS(binding.reduced_dof(binding.full_size()), std::out_of_range);
        CHECK_THROWS_AS(binding.supports(-1), std::out_of_range);
        CHECK_THROWS_AS(binding.supports(space.GetNE()), std::out_of_range);
        CHECK_THROWS_AS(binding.gather(mfem::Vector(binding.full_size() + 1)), std::invalid_argument);
        CHECK_THROWS_AS(binding.scatter(mfem::Vector(binding.size() + 1)), std::invalid_argument);
    }
}

TEST_CASE("MeanField supported field projections preserve canonical indexing and native GridFunctions", "[unit][discretization][meanfield]")
{
    mfem::Mesh mesh = split_mesh();
    mfem::H1_FECollection density_collection(blocks::density::mass::FunctionalSpace::order, 3);
    mfem::L2_FECollection velocity_collection(2, 3);
    mfem::H1_FECollection enthalpy_collection(3, 3);
    mfem::FiniteElementSpace density(&mesh, &density_collection);
    mfem::FiniteElementSpace enthalpy(&mesh, &enthalpy_collection);
    mfem::H1_FECollection potential_collection(blocks::gravity::potential::FunctionalSpace::order, 3);
    mfem::FiniteElementSpace potential(&mesh, &potential_collection);
    check_supported_projection<blocks::density::mass>(density);
    check_supported_projection<::serif::tests::equations::Enthalpy>(enthalpy);
    check_supported_projection<blocks::gravity::potential>(potential);
    for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM})
    {
        mfem::FiniteElementSpace displacement(&mesh, &enthalpy_collection, 3, ordering);
        check_supported_projection<blocks::displacement::vector>(displacement);
        mfem::FiniteElementSpace velocity(&mesh, &velocity_collection, 3, ordering);
        CHECK(velocity.GetFE(0)->GetOrder() == 2);
        CHECK(velocity.GetVDim() == 3);
        CHECK(velocity.GetOrdering() == ordering);
        check_supported_projection<DiscontinuousVelocity::vector>(velocity);
    }
    CHECK_THROWS_AS((backend::TermBinding<DiscontinuousVelocity::vector, Schema>{density}), std::invalid_argument);
}

TEST_CASE("MeanField stellar H1 fields retain interface values and exclude vacuum-only coefficients", "[unit][discretization][meanfield]")
{
    mfem::Mesh mesh = split_mesh();
    mfem::H1_FECollection h1(3, 3);
    mfem::FiniteElementSpace enthalpy_space(&mesh, &h1), density_space(&mesh, &h1);
    const backend::TermBinding<::serif::tests::equations::Enthalpy, Schema> enthalpy{enthalpy_space};
    const backend::TermBinding<blocks::density::mass, Schema> density{density_space};
    mfem::Array<int> stellar, vacuum;
    enthalpy_space.GetElementDofs(0, stellar);
    enthalpy_space.GetElementDofs(1, vacuum);
    int shared = 0;
    int excluded = 0;
    for (const int dof : vacuum)
    {
        const bool on_interface = std::find(stellar.begin(), stellar.end(), dof) != stellar.end();
        CHECK(enthalpy.reduced_dof(dof).has_value() == on_interface);
        shared += on_interface;
        excluded += !on_interface;
    }
    CHECK(shared > 0);
    CHECK(excluded > 0);
    density_space.GetElementDofs(0, stellar);
    density_space.GetElementDofs(1, vacuum);
    shared = 0;
    excluded = 0;
    for (const int dof : vacuum)
    {
        const bool on_interface = std::find(stellar.begin(), stellar.end(), dof) != stellar.end();
        CHECK(density.reduced_dof(dof).has_value() == on_interface);
        shared += on_interface;
        excluded += !on_interface;
    }
    CHECK(shared > 0);
    CHECK(excluded > 0);
    CHECK(enthalpy.supports(0));
    CHECK_FALSE(enthalpy.supports(1));
    CHECK(density.supports(0));
    CHECK_FALSE(density.supports(1));
}

TEST_CASE("MeanField field support follows alternate semantic material IDs", "[unit][discretization][meanfield]")
{
    mfem::Mesh mesh = split_mesh(17, 29);
    mfem::H1_FECollection h1(3, 3);
    mfem::FiniteElementSpace density(&mesh, &h1), enthalpy(&mesh, &h1);
    check_supported_projection<blocks::density::mass, AlternateSchema>(density);
    check_supported_projection<::serif::tests::equations::Enthalpy, AlternateSchema>(enthalpy);
}

TEST_CASE("MeanField supported maps reject stale bindings and malformed selections", "[unit][discretization][meanfield]")
{
    mfem::Mesh mesh = split_mesh();
    mfem::H1_FECollection collection(3, 3);
    mfem::FiniteElementSpace space(&mesh, &collection, 3);
    const backend::TermBinding<blocks::displacement::vector, Schema> binding{space};
    for (const std::vector<int>& entries : {std::vector<int>{-1, 2}, {1, binding.size()}, {1, 1, 3}, {1, 3, 2}})
    {
        mfem::Array<int> selected(static_cast<int>(entries.size()));
        std::copy(entries.begin(), entries.end(), selected.begin());
        CHECK_THROWS_AS(backend::DofSelection(binding, selected), std::invalid_argument);
    }
    mfem::Vector values(binding.size());
    values = 0.0;
    CHECK_NOTHROW(binding.validate_values(values));
    values[0] = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS_AS(binding.scatter(values), std::domain_error);
    mesh.GetElement(0)->SetAttribute(3);
    CHECK_THROWS(binding.validate());
}

TEST_CASE("MeanField registered spaces retain family order vector dimension and global scalar storage", "[unit][discretization][meanfield]")
{
    mfem::Mesh mesh = split_mesh();
    mfem::H1_FECollection h1(3, 3);
    mfem::H1_FECollection potential_collection(blocks::gravity::potential::FunctionalSpace::order, 3);
    mfem::FiniteElementSpace potential(&mesh, &potential_collection), displacement(&mesh, &h1, 3);
    const backend::TermBinding<blocks::gravity::potential, Schema> potential_binding{potential};
    const backend::TermBinding<blocks::displacement::vector, Schema> displacement_binding{displacement};
    CHECK(potential.GetFE(0)->GetOrder() == blocks::gravity::potential::FunctionalSpace::order);
    CHECK(displacement.GetFE(0)->GetOrder() == 3);
    CHECK(potential.GetVDim() == 1);
    CHECK(displacement.GetVDim() == 3);
    CHECK(potential_binding.size() == potential.GetVSize());
    CHECK(displacement_binding.size() == displacement.GetVSize());
    const backend::TermBinding<blocks::barotropic_constant::mass_normalization, Schema> global{mesh};
    CHECK(global.size() == 1);
    CHECK(global.full_size() == 1);
    CHECK_THROWS_AS(global.space(), std::logic_error);
    mfem::Vector scalar(1);
    scalar[0] = 2.7;
    CHECK(global.gather(global.scatter(scalar))[0] == 2.7);
    mfem::Vector local;
    global.element_values(0, scalar, local);
    CHECK(local[0] == 2.7);
    global.element_values(1, scalar, local);
    CHECK(local[0] == 0.0);
    CHECK_THROWS_AS((backend::TermBinding<blocks::density::mass, Schema>{displacement}), std::invalid_argument);
}

TEST_CASE("MeanField stellar boundary selections preserve exact reduced field ordering", "[unit][discretization][meanfield]")
{
    serif::tests::deformation::PairedMesh fixture(0, 0.0, 3, 4);
    const backend::TermBinding<::serif::tests::equations::Enthalpy, Schema> enthalpy{fixture.scalar_space};
    const backend::DofSelection selected = backend::select_boundary_dofs<domain::StellarSurfaceBoundary>(enthalpy);
    REQUIRE(selected.size() > 0);
    CHECK(selected.size() < enthalpy.size());
    CHECK(&selected.binding() == &enthalpy);
    mfem::Array<int> marker{1, 0};
    mfem::Array<int> full_boundary;
    fixture.scalar_space.GetEssentialTrueDofs(marker, full_boundary);
    REQUIRE(selected.size() == full_boundary.Size());
    for (int index = 0; index < full_boundary.Size(); ++index)
    {
        const std::optional<int> reduced = enthalpy.reduced_dof(full_boundary[index]);
        REQUIRE(reduced);
        CHECK(selected.reduced_dofs()[index] == *reduced);
        if (index > 0)
        {
            CHECK(selected.reduced_dofs()[index] > selected.reduced_dofs()[index - 1]);
        }
    }
    CHECK_NOTHROW(selected.validate());
}