// Port the exact reference mass-diagonal comparisons from MeanField normalization/stellar_equilibrium.cpp.
#include <algorithm>
#include <cmath>
#include <memory>
#include "mfem/mesh/ncmesh.hpp"

#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/bilinearform.hpp"
#include "mfem/fem/bilininteg.hpp"
#include "mfem/fem/gridfunc.hpp"
#include "mfem/general/forall.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/backend/mfem/discretization/mass_diagonal.hpp"
#include "serif/discretization/domain/schema/schemas.hpp"

namespace
{
    namespace discretization = serif::discretization;
    namespace backend = discretization::backend::mfem;
    namespace blocks = discretization::blocks;
    namespace domain = discretization::domain;
    namespace functional_space = discretization::functional_space;
    using Schema = domain::schema::CoreEnvelopeVacuumDomainSchema;

    /** @brief Exercise the same mass construction for stellar support and independent vector components. */
    struct TestField final : blocks::field
    {
        template <functional_space::SpaceDescriptor Space, int Order, std::size_t Rank, typename Domain = domain::StellarDomains>
        struct values final : blocks::term<TestField>
        {
            using FunctionalSpace = functional_space::DiscreteFunctionalSpace<Space, Order, Rank>;
            using Support = Domain;
        };
    };

    /** @brief Keep shared interface coefficients while excluding the second material's volume. */
    [[nodiscard]] mfem::Mesh reference_mesh(const mfem::Element::Type geometry)
    {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, geometry, 2.0, 1.5, 0.75);
        for (int element = 0; element < mesh.GetNE(); ++element)
        {
            mesh.GetElement(element)->SetAttribute(element < mesh.GetNE() / 2 ? 2 : 3);
        }
        mesh.SetAttributes();
        return mesh;
    }

    /** @brief Compare the production tensor shortcut or fallback against an independently assembled sparse matrix. */
    template <typename Term>
    void check_reference_diagonal(mfem::FiniteElementSpace& space)
    {
        const backend::TermBinding<Term, Schema> binding{space};
        const auto actual = backend::reference_mass_diagonal(binding);
        REQUIRE(actual);
        REQUIRE(actual->size() == static_cast<std::size_t>(binding.size()));

        mfem::Array<int> marker(space.GetMesh()->attributes.Max());
        for (int attribute = 1; attribute <= marker.Size(); ++attribute)
        {
            marker[attribute - 1] = Schema::template domain_id_belongs_to<typename Term::Support>(attribute);
        }
        mfem::BilinearForm reference(&space);
        if constexpr (blocks::SpatialTermTraits<Term>::rank == 1)
        {
            reference.AddDomainIntegrator(new mfem::VectorMassIntegrator(), marker);
        }
        else
        {
            reference.AddDomainIntegrator(new mfem::MassIntegrator(), marker);
        }
        reference.Assemble();
        reference.Finalize();
        mfem::Vector full_diagonal;
        reference.SpMat().GetDiag(full_diagonal);
        const mfem::Vector expected = binding.gather(full_diagonal);
        double error_squared = 0.0;
        for (int index = 0; index < expected.Size(); ++index)
        {
            CAPTURE(index, (*actual)[index], expected[index]);
            CHECK(std::isfinite((*actual)[index]));
            CHECK((*actual)[index] > 0.0);
            error_squared += std::pow((*actual)[index] - expected[index], 2);
        }
        CHECK(std::sqrt(error_squared) <= 2.0e-11 * expected.Norml2());
    }

    /** @brief Similar reference bodies change scalar volume metrics by length cubed and flux metrics by inverse length. */
    template <typename Term>
    void check_geometry_scaling(mfem::FiniteElementSpace& space, const double exponent)
    {
        const backend::TermBinding<Term, Schema> binding{space};
        const auto original = backend::reference_mass_diagonal(binding);
        REQUIRE(original);
        constexpr double length_scale = 2.5;
        mfem::VectorFunctionCoefficient dilation(3, [](const mfem::Vector& point, mfem::Vector& image)
        {
            image = point;
            image *= length_scale;
        });
        space.GetMesh()->Transform(dilation);
        const auto scaled = backend::reference_mass_diagonal(binding);
        REQUIRE(scaled);
        REQUIRE(scaled->size() == original->size());
        for (std::size_t index = 0; index < original->size(); ++index)
        {
            const double expected = (*original)[index] * std::pow(length_scale, exponent);
            CHECK(std::abs((*scaled)[index] - expected) <= 2.0e-12 * expected);
        }
    }
}

TEST_CASE("Reference mass diagonals match exact scalar assembly", "[unit][normalization][meanfield]")
{
    for (const mfem::Element::Type geometry : {mfem::Element::HEXAHEDRON, mfem::Element::TETRAHEDRON})
    {
        CAPTURE(geometry);
        mfem::Mesh mesh = reference_mesh(geometry);
        mfem::H1_FECollection h1(3, 3);
        mfem::FiniteElementSpace enthalpy(&mesh, &h1), density(&mesh, &h1);
        check_reference_diagonal<::serif::tests::equations::Enthalpy>(enthalpy);
        check_reference_diagonal<blocks::density::mass>(density);
        mfem::H1_FECollection potential_collection(blocks::gravity::potential::FunctionalSpace::order, 3);
        mfem::FiniteElementSpace potential(&mesh, &potential_collection);
        check_reference_diagonal<blocks::gravity::potential>(potential);
    }
}

TEST_CASE("Reference mass diagonals preserve vector component ordering", "[unit][normalization]")
{
    for (const mfem::Element::Type geometry : {mfem::Element::HEXAHEDRON, mfem::Element::TETRAHEDRON})
    {
        mfem::Mesh mesh = reference_mesh(geometry);
        mfem::H1_FECollection h1(3, 3);
        mfem::L2_FECollection l2(2, 3);
        for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM})
        {
            CAPTURE(geometry, ordering);
            mfem::FiniteElementSpace displacement(&mesh, &h1, 3, ordering), velocity(&mesh, &l2, 3, ordering);
            check_reference_diagonal<blocks::displacement::vector>(displacement);
            check_reference_diagonal<TestField::values<functional_space::L2, 2, 1>>(velocity);
        }
    }
}

TEST_CASE("Reference mass diagonals use curved geometry in both assembly paths", "[unit][normalization]")
{
    for (const mfem::Element::Type geometry : {mfem::Element::HEXAHEDRON, mfem::Element::TETRAHEDRON})
    {
        CAPTURE(geometry);
        mfem::Mesh mesh = reference_mesh(geometry);
        mesh.SetCurvature(2);
        mfem::VectorFunctionCoefficient curved(3, [](const mfem::Vector& point, mfem::Vector& image)
        {
            image = point;
            image[1] += 0.06 * point[0] * point[0];
            image[2] *= 1.0 + 0.04 * point[0];
        });
        mesh.Transform(curved);
        mfem::H1_FECollection h1(3, 3);
        mfem::FiniteElementSpace enthalpy(&mesh, &h1), density(&mesh, &h1), displacement(&mesh, &h1, 3);
        check_reference_diagonal<::serif::tests::equations::Enthalpy>(enthalpy);
        check_reference_diagonal<blocks::density::mass>(density);
        check_reference_diagonal<blocks::displacement::vector>(displacement);
    }
}

TEST_CASE("Reference metrics retain scalar length scaling", "[unit][normalization]")
{
    SECTION("Scalar volume coefficients")
    {
        mfem::Mesh mesh = reference_mesh(mfem::Element::HEXAHEDRON);
        mfem::H1_FECollection h1(3, 3);
        mfem::FiniteElementSpace enthalpy(&mesh, &h1);
        check_geometry_scaling<::serif::tests::equations::Enthalpy>(enthalpy, 3.0);
    }
}

TEST_CASE("Reference metrics retain supported ordering and reject stale material bindings", "[unit][normalization]")
{
    mfem::Mesh mesh = reference_mesh(mfem::Element::HEXAHEDRON);
    mfem::H1_FECollection h1(3, 3);
    mfem::FiniteElementSpace space(&mesh, &h1);
    const backend::TermBinding<::serif::tests::equations::Enthalpy, Schema> binding{space};
    const auto diagonal = backend::reference_mass_diagonal(binding);
    REQUIRE(diagonal);
    CHECK(diagonal->size() < static_cast<std::size_t>(space.GetVSize()));

    mesh.GetElement(0)->SetAttribute(3);
    CHECK_THROWS(backend::reference_mass_diagonal(binding));
    mesh.SetAttributes();
    const backend::TermBinding<::serif::tests::equations::Enthalpy, Schema> empty_binding{space};
    const auto empty = backend::reference_mass_diagonal(empty_binding);
    REQUIRE(empty);
    CHECK(empty->empty());
}

TEST_CASE("Reference metrics reject zero volume rather than constructing singular normalization", "[unit][normalization]")
{
    mfem::Mesh mesh = reference_mesh(mfem::Element::HEXAHEDRON);
    mfem::VectorFunctionCoefficient collapsed(3, [](const mfem::Vector& point, mfem::Vector& image)
    {
        image = point;
        image[2] = 0.0;
    });
    mesh.Transform(collapsed);
    mfem::H1_FECollection h1(3, 3);
    mfem::FiniteElementSpace space(&mesh, &h1);
    const backend::TermBinding<::serif::tests::equations::Enthalpy, Schema> binding{space};
    const auto diagonal = backend::reference_mass_diagonal(binding);
    REQUIRE_FALSE(diagonal);
    CHECK(diagonal.error().find("positive finite") != std::string::npos);
}

namespace {
    /** @brief Assemble an independent dense element oracle before MFEM applies hanging-node constraints. */
    class ReconstructedMass final : public mfem::BilinearFormIntegrator {
    public:
        /** @brief Borrow the weighted binding and select a verification quadrature increment. */
        ReconstructedMass(const backend::DofMap& binding, int increment) : m_binding(binding), m_increment(increment) {}
        /** @brief Preserve every local cross term, so the oracle tests constrained diagonal reduction. */
        void AssembleElementMatrix(const mfem::FiniteElement& element, mfem::ElementTransformation& transformation,
            mfem::DenseMatrix& matrix) override {
            const auto* reconstruction = m_binding.reconstruction_element(transformation.ElementNo);
            const int order = mfem::MassIntegrator::GetRule(element, element, transformation).GetOrder() + m_increment;
            const auto& rule = mfem::IntRules.Get(element.GetGeomType(), order);
            matrix.SetSize(element.GetDof());
            matrix = 0.0;
            mfem::Vector shapes(element.GetDof());
            mfem::DenseMatrix derivatives;
            for (int index = 0; index < rule.GetNPoints(); ++index) {
                const auto& point = rule.IntPoint(index);
                if (reconstruction != nullptr) { reconstruction->sample(element, point, shapes, derivatives); }
                else { element.CalcShape(point, shapes); }
                transformation.SetIntPoint(&point);
                const double weight = point.weight * transformation.Weight();
                for (int row = 0; row < element.GetDof(); ++row) {
                    for (int column = 0; column < element.GetDof(); ++column) { matrix(row, column) += weight * shapes[row] * shapes[column]; }
                }
            }
        }
    private:
        const backend::DofMap& m_binding;
        int m_increment;
    };

    /** @brief Form the actual sparse mass matrix, then reduce all couplings before reading its diagonal. */
    [[nodiscard]] mfem::Vector reconstructed_diagonal_oracle(const backend::DofMap& binding, int increment) {
        mfem::BilinearForm mass(&binding.space());
        mass.AddDomainIntegrator(new ReconstructedMass(binding, increment));
        mass.Assemble();
        mass.Finalize();
        mfem::Vector diagonal;
        if (binding.constrained()) {
            const auto& prolongation = *binding.space().GetConformingProlongation();
            const std::unique_ptr<mfem::SparseMatrix> reduced(mfem::RAP(prolongation, mass.SpMat(), prolongation));
            reduced->GetDiag(diagonal);
            return binding.gather_true(diagonal);
        }
        mass.SpMat().GetDiag(diagonal);
        return binding.gather(diagonal);
    }
}

TEST_CASE("Surface-power reference metrics retain constrained cross terms and converge under quadrature", "[unit][normalization][reconstruction]") {
    for (bool constrained : {false, true}) {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON);
        mfem::Array<int> vertices;
        for (int boundary = 0; boundary < mesh.GetNBE(); ++boundary) {
            mesh.GetBdrElementVertices(boundary, vertices);
            bool surface = true;
            for (int vertex : vertices) { surface = surface && mesh.GetVertex(vertex)[2] == 0.0; }
            mesh.GetBdrElement(boundary)->SetAttribute(surface ? 1 : 2);
        }
        mesh.SetAttributes();
        if (constrained) {
            mfem::Array<mfem::Refinement> refinement;
            refinement.Append(mfem::Refinement(0, mfem::Refinement::Y));
            mesh.GeneralRefinement(refinement, 1);
        }
        mfem::H1_FECollection collection(3, 3);
        mfem::FiniteElementSpace space(&mesh, &collection);
        const backend::TermBinding<blocks::density::mass, Schema> ordinary(space);
        const auto plan = std::make_shared<backend::ReconstructionPlan>(ordinary, 1,
            std::make_shared<discretization::reconstruction::SurfacePowerProfile>(5.0 / 3.0));
        const backend::TermBinding<blocks::density::mass, Schema> binding(space, plan);
        REQUIRE(binding.constrained() == constrained);
        const auto actual = backend::reference_mass_diagonal(binding);
        REQUIRE(actual);
        const mfem::Vector oracle = reconstructed_diagonal_oracle(binding, plan->integration_order_increment());
        const mfem::Vector refined = reconstructed_diagonal_oracle(binding, plan->integration_order_increment() + 16);
        for (int index = 0; index < binding.size(); ++index) {
            CAPTURE(constrained, index);
            CHECK((*actual)[index] > 0.0);
            CHECK(std::abs((*actual)[index] - oracle[index]) <= 2e-11 * oracle[index]);
            // Fractional powers are not integrated exactly by a finite Gaussian rule.
            CHECK(std::abs((*actual)[index] - refined[index]) <= 3e-4 * refined[index]);
        }
    }
}