#include <cmath>

#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/coefficient.hpp"
#include "mfem/mesh/mesh.hpp"
#include "serif/backend/mfem/projection/radial.hpp"

namespace {
    namespace projection = serif::utils::projection::backend::mfem;

    /** @brief Refine one of two adjacent cells so the shared edge has constrained child coefficients. */
    [[nodiscard]] mfem::Mesh nonconforming_mesh() {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian2D(2, 1, mfem::Element::QUADRILATERAL, true, 2.0, 1.0);
        mesh.EnsureNCMesh();
        mfem::Array<int> selected(1);
        selected[0] = 0;
        mesh.GeneralRefinement(selected);
        return mesh;
    }

    /** @brief Independently measure the distance from the true finite-element subspace through its constraint matrices. */
    [[nodiscard]] double constraint_defect(const mfem::GridFunction& field) {
        const mfem::FiniteElementSpace& space = *field.FESpace();
        const mfem::SparseMatrix* restriction = space.GetConformingRestriction();
        const mfem::SparseMatrix* prolongation = space.GetConformingProlongation();
        REQUIRE(restriction != nullptr);
        REQUIRE(prolongation != nullptr);
        REQUIRE(space.GetTrueVSize() < space.GetVSize());
        mfem::Vector true_values(restriction->Height()), constrained(space.GetVSize());
        restriction->Mult(field, true_values);
        prolongation->Mult(true_values, constrained);
        constrained -= field;
        return constrained.Norml2();
    }
}

TEST_CASE("Radial scalar projection satisfies hanging-node constraints after nonlinear sampling", "[projection][nonconforming]") {
    mfem::Mesh mesh = nonconforming_mesh();
    mfem::H1_FECollection collection(1, 2);
    mfem::FiniteElementSpace space(&mesh, &collection);
    mfem::FunctionCoefficient coefficient([](const mfem::Vector& position) {
        return position * position;
    });
    mfem::GridFunction native(&space);
    native.ProjectCoefficient(coefficient);
    REQUIRE(constraint_defect(native) > 0.1);
    const auto projected = projection::project_scalar(space, [](double radius) {
        return radius * radius;
    });
    REQUIRE(projected);
    CHECK(constraint_defect(*projected) < 1.0e-13);
}

TEST_CASE("Radial vector projection satisfies hanging-node and component constraints", "[projection][nonconforming]") {
    mfem::Mesh mesh = nonconforming_mesh();
    mfem::H1_FECollection vector_collection(1, 2);
    mfem::VectorFunctionCoefficient coefficient(2, [](const mfem::Vector& position, mfem::Vector& value) {
        value = position;
        value *= position * position;
    });
    const auto check = [&](mfem::FiniteElementSpace& space) {
        mfem::GridFunction native(&space);
        native.ProjectCoefficient(coefficient);
        REQUIRE(constraint_defect(native) > 1.0e-3);
        const auto projected = projection::project_vector(space, [](double radius) {
            return radius * radius * radius;
        });
        REQUIRE(projected);
        CHECK(constraint_defect(*projected) < 1.0e-13);
    };
    for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM}) {
        mfem::FiniteElementSpace vector(&mesh, &vector_collection, 2, ordering);
        check(vector);
    }
}

TEST_CASE("Constrained radial interpolation retains representable scalar fields and exact vector divergence", "[projection][nonconforming][math]") {
    mfem::Mesh mesh = nonconforming_mesh();
    mfem::H1_FECollection scalar_collection(2, 2);
    mfem::H1_FECollection vector_collection(1, 2);
    mfem::FiniteElementSpace scalar_space(&mesh, &scalar_collection), flux_space(&mesh, &vector_collection, 2);
    const auto scalar = projection::project_scalar(scalar_space, [](double radius) {
        return radius * radius;
    });
    const auto flux = projection::project_vector(flux_space, [](double radius) {
        return radius;
    });
    REQUIRE(scalar);
    REQUIRE(flux);
    mfem::FunctionCoefficient exact_scalar([](const mfem::Vector& position) {
        return position * position;
    });
    mfem::VectorFunctionCoefficient exact_vector(2, [](const mfem::Vector& position, mfem::Vector& value) {
        value = position;
    });
    mfem::ConstantCoefficient divergence(2.0);
    CHECK(scalar->ComputeL2Error(exact_scalar) < 2.0e-13);
    CHECK(flux->ComputeL2Error(exact_vector) < 2.0e-13);
    CHECK(flux->ComputeDivError(&divergence) < 2.0e-13);
    CHECK(constraint_defect(*scalar) < 1.0e-13);
    CHECK(constraint_defect(*flux) < 1.0e-13);
}