#include <algorithm>
#include <array>
#include <algorithm>
#include <cmath>
#include <memory>
#include <limits>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace {
    using ScalarTable = mean_field::fem::ScalarReferenceTable;
    using VectorTable = mean_field::fem::VectorReferenceTable;
    using TableCache  = mean_field::fem::ReferenceTableCache;

    static_assert(std::is_same_v<
                  decltype(std::declval<const TableCache &>().GetScalarTable(
                      std::declval<const mfem::FiniteElement &>(),
                      std::declval<const mfem::IntegrationRule &>()
                  )),
                  std::shared_ptr<const ScalarTable>>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const ScalarTable &>().GetValues()),
                  const mfem::DenseMatrix &>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const ScalarTable &>().GetGradients(0)),
                  const mfem::DenseMatrix &>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const TableCache &>().GetVectorTable(
                      std::declval<const mfem::FiniteElement &>(),
                      std::declval<const mfem::IntegrationRule &>()
                  )),
                  std::shared_ptr<const VectorTable>>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const VectorTable &>().GetValues(0)),
                  const mfem::DenseMatrix &>);

    void CheckMatrixExactly(
        const mfem::DenseMatrix &actual,
        const mfem::DenseMatrix &expected
    ) {
        REQUIRE(actual.Height() == expected.Height());
        REQUIRE(actual.Width() == expected.Width());
        for (int column = 0; column < actual.Width(); ++column) {
            for (int row = 0; row < actual.Height(); ++row) {
                CHECK(actual(row, column) == expected(row, column));
            }
        }
    }

    void CheckScalarTable(
        const ScalarTable &table,
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    ) {
        REQUIRE(table.GetPointCount() == rule.GetNPoints());
        REQUIRE(table.GetDofCount() == element.GetDof());
        REQUIRE(table.GetDimension() == element.GetDim());
        REQUIRE(table.GetValues().Height() == rule.GetNPoints());
        REQUIRE(table.GetValues().Width() == element.GetDof());

        mfem::Vector shape(element.GetDof());
        mfem::DenseMatrix gradient(element.GetDof(), element.GetDim());
        for (int point = 0; point < rule.GetNPoints(); ++point) {
            element.CalcShape(rule.IntPoint(point), shape);
            element.CalcDShape(rule.IntPoint(point), gradient);
            for (int dof = 0; dof < element.GetDof(); ++dof) {
                CHECK(table.GetValues()(point, dof) == shape(dof));
            }
            CheckMatrixExactly(table.GetGradients(point), gradient);
        }
    }

    void CheckVectorTable(
        const VectorTable &table,
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    ) {
        REQUIRE(table.GetPointCount() == rule.GetNPoints());
        REQUIRE(table.GetDofCount() == element.GetDof());
        REQUIRE(table.GetDimension() == element.GetRangeDim());
        mfem::DenseMatrix shape(element.GetDof(), element.GetRangeDim());
        for (int point = 0; point < rule.GetNPoints(); ++point) {
            element.CalcVShape(rule.IntPoint(point), shape);
            CheckMatrixExactly(table.GetValues(point), shape);
        }
    }

    mfem::IntegrationRule CopyRule(const mfem::IntegrationRule &source) {
        mfem::IntegrationRule copy(source.GetNPoints());
        copy.SetOrder(source.GetOrder());
        for (int point = 0; point < source.GetNPoints(); ++point) {
            copy.IntPoint(point) = source.IntPoint(point);
        }
        return copy;
    }
} // namespace

TEST_CASE(
    "Reference Table Cache Matches Scalar MFEM Values And Gradients",
    tags::unit &tags::quadrature
) {
    for (const int dimension : std::array{2, 3}) {
        const auto geometry = dimension == 2 ? mfem::Geometry::SQUARE : mfem::Geometry::CUBE;
        for (const int order : std::array{1, 3}) {
            CAPTURE(dimension, order);
            mfem::H1_FECollection h1(order, dimension);
            mfem::L2_FECollection l2(order - 1, dimension);
            TableCache cache;
            const mfem::IntegrationRule &rule = mfem::IntRules.Get(geometry, 2 * order + 1);
            for (const mfem::FiniteElement *element :
                 std::array{h1.FiniteElementForGeometry(geometry), l2.FiniteElementForGeometry(geometry)}) {
                REQUIRE(element != nullptr);
                const auto table = cache.GetScalarTable(*element, rule);
                REQUIRE(table != nullptr);
                CheckScalarTable(*table, *element, rule);
            }
        }
    }
}

TEST_CASE(
    "Reference Table Cache Matches RT Reference Values And Shares Equal Rules",
    tags::unit &tags::quadrature
) {
    for (const int dimension : std::array{2, 3}) {
        const auto geometry = dimension == 2 ? mfem::Geometry::SQUARE : mfem::Geometry::CUBE;
        for (const int order : std::array{0, 2}) {
            CAPTURE(dimension, order);
            mfem::RT_FECollection standard(order, dimension);
            mfem::RT_FECollection integrated(
                order, dimension, mfem::BasisType::GaussLobatto, mfem::BasisType::IntegratedGLL
            );
            const mfem::FiniteElement &standardElement   = *standard.FiniteElementForGeometry(geometry);
            const mfem::FiniteElement &integratedElement = *integrated.FiniteElementForGeometry(geometry);
            const mfem::IntegrationRule &rule            = mfem::IntRules.Get(geometry, 2 * order + 3);
            const mfem::IntegrationRule copiedRule       = CopyRule(rule);
            const TableCache cache;
            const auto standardTable   = cache.GetVectorTable(standardElement, rule);
            const auto integratedTable = cache.GetVectorTable(integratedElement, rule);
            REQUIRE(standardTable != nullptr);
            REQUIRE(integratedTable != nullptr);
            CHECK(cache.GetVectorTable(standardElement, copiedRule).get() == standardTable.get());
            CHECK(cache.GetVectorTable(integratedElement, copiedRule).get() == integratedTable.get());
            CHECK(standardTable.get() != integratedTable.get());
            CheckVectorTable(*standardTable, standardElement, rule);
            CheckVectorTable(*integratedTable, integratedElement, rule);
        }
    }
}

TEST_CASE(
    "Reference Table Cache Shares Equal Rules And Distinguishes Rule Contents",
    tags::unit &tags::quadrature
) {
    mfem::H1_FECollection collection(3, 3);
    const mfem::FiniteElement &element      = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
    const mfem::IntegrationRule &rule       = mfem::IntRules.Get(mfem::Geometry::CUBE, 7);
    mfem::IntegrationRule copiedRule        = CopyRule(rule);
    mfem::IntegrationRule movedPointRule    = CopyRule(rule);
    mfem::IntegrationRule changedWeightRule = CopyRule(rule);
    movedPointRule.IntPoint(0).x += 0.03125;
    changedWeightRule.IntPoint(0).weight *= 1.25;

    const TableCache cache;
    const auto original = cache.GetScalarTable(element, rule);
    const mfem::DenseMatrix originalValues(original->GetValues());
    const auto copy          = cache.GetScalarTable(element, copiedRule);
    const auto movedPoint    = cache.GetScalarTable(element, movedPointRule);
    const auto changedWeight = cache.GetScalarTable(element, changedWeightRule);

    CHECK(copy.get() == original.get());
    CHECK(movedPointRule.GetOrder() == rule.GetOrder());
    CHECK(changedWeightRule.GetOrder() == rule.GetOrder());
    CHECK(movedPoint.get() != original.get());
    CHECK(changedWeight.get() != original.get());
    CHECK(changedWeight.get() != movedPoint.get());
    CheckScalarTable(*movedPoint, element, movedPointRule);
    CheckScalarTable(*changedWeight, element, changedWeightRule);
    CheckMatrixExactly(original->GetValues(), originalValues);

    // Rule identity is its contents, not its address, including after mutation.
    copiedRule.IntPoint(0).x = movedPointRule.IntPoint(0).x;
    CHECK(cache.GetScalarTable(element, copiedRule).get() == movedPoint.get());
    CHECK(cache.GetScalarTable(element, rule).get() == original.get());
}

TEST_CASE(
    "Reference Table Cache Distinguishes Scalar Basis Variants",
    tags::unit &tags::quadrature
) {
    constexpr int dimension = 3;
    constexpr int order     = 3;
    mfem::H1_FECollection nodalH1(order, dimension, mfem::BasisType::GaussLobatto);
    mfem::H1_FECollection positiveH1(order, dimension, mfem::BasisType::Positive);
    mfem::L2_FECollection openL2(order, dimension, mfem::BasisType::GaussLegendre);
    mfem::L2_FECollection closedL2(order, dimension, mfem::BasisType::GaussLobatto);
    const mfem::IntegrationRule &rule = mfem::IntRules.Get(mfem::Geometry::CUBE, 5);
    const TableCache cache;
    std::array<std::shared_ptr<const ScalarTable>, 4> tables;
    const std::array<const mfem::FiniteElement *, 4> elements{
        nodalH1.FiniteElementForGeometry(mfem::Geometry::CUBE),
        positiveH1.FiniteElementForGeometry(mfem::Geometry::CUBE),
        openL2.FiniteElementForGeometry(mfem::Geometry::CUBE), closedL2.FiniteElementForGeometry(mfem::Geometry::CUBE)
    };
    for (std::size_t index = 0; index < elements.size(); ++index) {
        REQUIRE(elements[index] != nullptr);
        REQUIRE(elements[index]->GetOrder() == order);
        REQUIRE(elements[index]->GetDof() == elements[0]->GetDof());
        tables[index] = cache.GetScalarTable(*elements[index], rule);
        CheckScalarTable(*tables[index], *elements[index], rule);
        for (std::size_t previous = 0; previous < index; ++previous) {
            CHECK(tables[index].get() != tables[previous].get());
        }
    }
}

TEST_CASE(
    "Reference Table Cache Published Scalar Storage Outlives Its Cache",
    tags::unit &tags::quadrature
) {
    std::shared_ptr<const ScalarTable> retained;
    mfem::DenseMatrix expectedValues;
    mfem::DenseMatrix expectedGradient;
    {
        // FE objects remain immutable and alive throughout the cache lifetime.
        mfem::H1_FECollection collection(3, 2);
        const mfem::FiniteElement &element = *collection.FiniteElementForGeometry(mfem::Geometry::SQUARE);
        const mfem::IntegrationRule rule   = CopyRule(mfem::IntRules.Get(mfem::Geometry::SQUARE, 7));
        const TableCache cache;
        retained = cache.GetScalarTable(element, rule);
        CheckScalarTable(*retained, element, rule);
        expectedValues   = retained->GetValues();
        expectedGradient = retained->GetGradients(0);
    }
    REQUIRE(retained != nullptr);
    CheckMatrixExactly(retained->GetValues(), expectedValues);
    CheckMatrixExactly(retained->GetGradients(0), expectedGradient);
}

TEST_CASE(
    "Reference Table Cache Published RT Storage Outlives Its Cache",
    tags::unit &tags::quadrature
) {
    std::shared_ptr<const VectorTable> retained;
    mfem::DenseMatrix firstExpected;
    mfem::DenseMatrix lastExpected;
    {
        mfem::RT_FECollection collection(2, 3, mfem::BasisType::GaussLobatto, mfem::BasisType::IntegratedGLL);
        const mfem::FiniteElement &element = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
        const mfem::IntegrationRule rule   = CopyRule(mfem::IntRules.Get(mfem::Geometry::CUBE, 7));
        const TableCache cache;
        retained = cache.GetVectorTable(element, rule);
        CheckVectorTable(*retained, element, rule);
        firstExpected = retained->GetValues(0);
        lastExpected  = retained->GetValues(rule.GetNPoints() - 1);
    }
    REQUIRE(retained != nullptr);
    CheckMatrixExactly(retained->GetValues(0), firstExpected);
    CheckMatrixExactly(retained->GetValues(retained->GetPointCount() - 1), lastExpected);
}

TEST_CASE("Tensor reference evaluation matches native H1 hex values and gradients",
          tags::unit &tags::quadrature) {
    for (const int order : {1, 2, 4}) {
        for (const int basis : {mfem::BasisType::GaussLobatto, mfem::BasisType::ClosedUniform}) {
            mfem::H1_FECollection collection(order, 3, basis);
            const auto &element = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
            const auto rule = CopyRule(mfem::IntRules.Get(mfem::Geometry::CUBE, 2 * order + 3));
            const TableCache cache;
            const auto table = cache.GetScalarTable(element, rule);
            REQUIRE(table->HasTensorEvaluation3D());
            mfem::DenseMatrix dofs(element.GetDof(), 3), actual;
            mfem::Vector workspace;
            for (int c = 0; c < 3; ++c) for (int i = 0; i < element.GetDof(); ++i)
                dofs(i, c) = std::sin(0.37 * (i + 1) + c) + 0.21 * std::cos(0.13 * i - c);
            table->EvaluateTensor3D(dofs, actual, workspace);
            REQUIRE(actual.Height() == rule.GetNPoints());
            REQUIRE(actual.Width() == 12);
            for (int q = 0; q < rule.GetNPoints(); ++q) for (int c = 0; c < 3; ++c) {
                for (int derivative = 0; derivative < 4; ++derivative) {
                    double expected = 0.0;
                    for (int i = 0; i < element.GetDof(); ++i)
                        expected += dofs(i, c) * (derivative == 0 ? table->GetValues()(q, i) :
                            table->GetGradients(q)(i, derivative - 1));
                    CHECK(std::abs(actual(q, 4 * c + derivative) - expected) <
                          2.0e-12 * std::max(1.0, std::abs(expected)));
                }
            }
            // Reused caller scratch must not retain a previous direction.
            dofs = 0.0;
            table->EvaluateTensor3D(dofs, actual, workspace);
            for (int c = 0; c < actual.Width(); ++c) for (int q = 0; q < actual.Height(); ++q)
                CHECK(actual(q, c) == 0.0);
        }
    }
}

TEST_CASE("Tensor reference evaluation rejects non-Cartesian rules and owns its data",
          tags::unit &tags::quadrature) {
    std::shared_ptr<const ScalarTable> retained;
    mfem::DenseMatrix input, expected;
    {
        mfem::H1_FECollection collection(3, 3);
        const auto &element = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
        auto rule = CopyRule(mfem::IntRules.Get(mfem::Geometry::CUBE, 7));
        const TableCache cache;
        retained = cache.GetScalarTable(element, rule);
        REQUIRE(retained->HasTensorEvaluation3D());
        input.SetSize(element.GetDof(), 1);
        input = 1.0;
        mfem::Vector scratch;
        retained->EvaluateTensor3D(input, expected, scratch);
        rule.IntPoint(1).y += 0.001;
        CHECK_FALSE(cache.GetScalarTable(element, rule)->HasTensorEvaluation3D());
        const auto &tet = *collection.FiniteElementForGeometry(mfem::Geometry::TETRAHEDRON);
        CHECK_FALSE(cache.GetScalarTable(tet, mfem::IntRules.Get(mfem::Geometry::TETRAHEDRON, 7))
                        ->HasTensorEvaluation3D());
    }
    mfem::DenseMatrix actual;
    mfem::Vector scratch;
    retained->EvaluateTensor3D(input, actual, scratch);
    CheckMatrixExactly(actual, expected);
}

TEST_CASE("Tensor Reference Transposes Match Dense Assembly And Adjoint Products", tags::unit &tags::quadrature) {
    for (int order : {1, 2, 4}) {
        mfem::H1_FECollection h1(order, 3);
        const auto &fe = *h1.FiniteElementForGeometry(mfem::Geometry::CUBE);
        const auto &rule = mfem::IntRules.Get(mfem::Geometry::CUBE, 2*order+3);
        TableCache cache;
        auto table = cache.GetScalarTable(fe, rule);
        REQUIRE(table->HasTensorEvaluation3D());
        mfem::DenseMatrix dual(rule.GetNPoints(), 3), actual;
        for (int c = 0; c < 3; ++c) for (int q = 0; q < dual.Height(); ++q)
            dual(q,c) = std::sin(0.17*(q+1)*(c+1));
        mfem::Vector workspace, expected(fe.GetDof());
        for (int repeat = 0; repeat < 2; ++repeat) {
            if (repeat) dual = 0.0;
            table->ApplyTransposeTensor3D(dual, actual, workspace);
            for (int c = 0; c < 3; ++c) {
                table->GetValues().MultTranspose(dual.GetColumn(c), expected.GetData());
                for (int i = 0; i < fe.GetDof(); ++i)
                    CHECK(std::abs(actual(i,c)-expected(i)) <= 2.0e-12 * std::max(1.0, expected.Normlinf()));
            }
        }
    }
    for (int order : {0, 1, 3}) for (int basis : {mfem::BasisType::GaussLegendre, mfem::BasisType::IntegratedGLL}) {
        CAPTURE(order, basis);
        mfem::RT_FECollection rt(order, 3, mfem::BasisType::GaussLobatto, basis);
        const auto &fe = *rt.FiniteElementForGeometry(mfem::Geometry::CUBE);
        auto rule = CopyRule(mfem::IntRules.Get(mfem::Geometry::CUBE, 2*order+5));
        TableCache cache;
        auto table = cache.GetVectorTable(fe, rule);
        REQUIRE(table->HasTensorEvaluation3D());
        mfem::Vector dofs(fe.GetDof()), assembled, expected(fe.GetDof()), point(3), workspace;
        mfem::DenseMatrix values, dual(rule.GetNPoints(), 3);
        for (int i = 0; i < dofs.Size(); ++i) dofs(i) = std::sin(0.53*(i+1));
        for (int c = 0; c < 3; ++c) for (int q = 0; q < dual.Height(); ++q) dual(q,c) = std::cos(0.31*(q+1)+c);
        for (int repeat = 0; repeat < 2; ++repeat) {
            if (repeat) { dofs = 0.0; dual = 0.0; }
            table->EvaluateTensor3D(dofs, values, workspace);
            expected = 0.0;
            double primalProduct = 0.0;
            for (int q = 0; q < rule.GetNPoints(); ++q) {
                table->GetValues(q).MultTranspose(dofs, point);
                for (int c = 0; c < 3; ++c) {
                    CHECK(std::abs(values(q,c)-point(c)) <= 2.0e-11 * std::max(1.0, point.Normlinf()));
                    primalProduct += values(q,c)*dual(q,c);
                    point(c) = dual(q,c);
                }
                table->GetValues(q).AddMult(point, expected);
            }
            table->ApplyTransposeTensor3D(dual, assembled, workspace);
            mfem::Vector difference(assembled); difference -= expected;
            CHECK(difference.Norml2() <= 2.0e-11 * std::max(1.0, expected.Norml2()));
            CHECK(std::abs(primalProduct - dofs*assembled) <= 2.0e-11 * std::max(1.0, std::abs(primalProduct)));
        }
        rule.IntPoint(rule.GetNPoints()-1).x += 0.017;
        CHECK_FALSE(cache.GetVectorTable(fe, rule)->HasTensorEvaluation3D());
    }
}

TEST_CASE("RT Tensor Plans Own Their Factors After Cache And FE Destruction", tags::unit &tags::quadrature) {
    std::shared_ptr<const VectorTable> table;
    mfem::Vector dofs, expected;
    mfem::DenseMatrix values, dual;
    {
        mfem::RT_FECollection rt(2, 3);
        TableCache cache;
        auto rule = CopyRule(mfem::IntRules.Get(mfem::Geometry::CUBE, 9));
        const auto &fe = *rt.FiniteElementForGeometry(mfem::Geometry::CUBE);
        table = cache.GetVectorTable(fe, rule);
        REQUIRE(table->HasTensorEvaluation3D());
        dofs.SetSize(fe.GetDof()); dofs = 0.3;
        dual.SetSize(rule.GetNPoints(), 3); dual = 0.17;
        mfem::Vector workspace;
        table->ApplyTransposeTensor3D(dual, expected, workspace);
    }
    mfem::Vector workspace, actual;
    table->EvaluateTensor3D(dofs, values, workspace);
    table->ApplyTransposeTensor3D(dual, actual, workspace);
    actual -= expected;
    CHECK(actual.Norml2() == 0.0);
}


TEST_CASE("L2 Tensor Values And Transposes Match Dense Reference", tags::unit &tags::quadrature) {
    for (int order : {0, 1, 3}) for (int basis : {mfem::BasisType::GaussLegendre, mfem::BasisType::OpenUniform}) {
        mfem::L2_FECollection collection(order, 3, basis);
        const auto &element = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
        const auto &rule = mfem::IntRules.Get(mfem::Geometry::CUBE, 2*order+3);
        TableCache cache;
        const auto table = cache.GetScalarTable(element, rule);
        REQUIRE(table->HasTensorEvaluation3D());
        mfem::Vector dofs(element.GetDof()), dual(rule.GetNPoints()), values, expected(rule.GetNPoints()),
                     transpose, expectedTranspose(element.GetDof()), workspace;
        for (int i = 0; i < dofs.Size(); ++i) dofs(i) = std::sin(0.37*(i+1));
        for (int q = 0; q < dual.Size(); ++q) dual(q) = std::cos(0.19*(q+1));
        table->EvaluateValuesTensor3D(dofs, values, workspace);
        table->GetValues().Mult(dofs, expected);
        mfem::Vector difference(values); difference -= expected;
        CHECK(difference.Norml2() <= 2.0e-12*std::max(1.0, expected.Norml2()));
        table->ApplyTransposeTensor3D(dual, transpose, workspace);
        table->GetValues().MultTranspose(dual, expectedTranspose);
        difference = transpose; difference -= expectedTranspose;
        CHECK(difference.Norml2() <= 2.0e-12*std::max(1.0, expectedTranspose.Norml2()));
        CHECK(std::abs((values*dual) - (dofs*transpose)) <= 2.0e-12*std::max(1.0, std::abs(values*dual)));
        dofs = 0.0; dual = 0.0;
        table->EvaluateValuesTensor3D(dofs, values, workspace);
        table->ApplyTransposeTensor3D(dual, transpose, workspace);
        CHECK(values.Norml2() == 0.0);
        CHECK(transpose.Norml2() == 0.0);
    }
}

TEST_CASE("Geometry Inspector Reuses And Invalidates Fixed Geometry", tags::unit &tags::geometry) {
    for (const auto ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM}) {
        auto serial = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON);
        serial.SetCurvature(2);
        mfem::ParMesh mesh(MPI_COMM_WORLD, serial);
        mfem::H1_FECollection collection(2, 3);
        mfem::ParFiniteElementSpace space(&mesh, &collection, 3, ordering);
        mfem::ParGridFunction field(&space);
        mfem::VectorFunctionCoefficient displacement(3, [](const mfem::Vector &x, mfem::Vector &u) {
            u.SetSize(3);
            u(0) = 0.08*x(0) + 0.02*x(1)*x(2);
            u(1) = -0.03*x(1) + 0.01*x(0)*x(2);
            u(2) = 0.04*x(2);
        });
        field.ProjectCoefficient(displacement);
        mfem::Vector state; field.GetTrueDofs(state);
        mean_field::fem::PreparedGeometryInspector inspector;
        auto compare = [&] {
            const double fast = inspector.Inspect(space, state);
            const double reference = inspector.Inspect(space, state, true);
            CHECK(std::isfinite(fast));
            CHECK(std::abs(fast-reference) <= 2.0e-12*std::max(1.0, std::abs(reference)));
            return fast;
        };
        compare();
        CHECK(inspector.GetTensorElementCount() == mesh.GetNE());
        const auto builds = inspector.GetGeometryBuildCount();
        state *= 0.5;
        const double before = compare();
        CHECK(inspector.GetGeometryBuildCount() == builds);

        // Same FE/mesh topology, changed curved computational coordinates.
        mfem::VectorFunctionCoefficient stretch(3, [](const mfem::Vector &x, mfem::Vector &y) {
            y = x; y(0) = 1.2*x(0); y(1) += 0.04*x(0)*x(2);
        });
        mesh.Transform(stretch);
        const double after = compare();
        CHECK(std::abs(after-before) > 1.0e-6);
        CHECK(inspector.GetGeometryBuildCount() == builds + mesh.GetNE());
        inspector.ClearCache();
        CHECK(std::abs(compare()-after) < 2.0e-14);
        CHECK(inspector.GetGeometryBuildCount() == builds + 2*mesh.GetNE());

        mfem::VectorFunctionCoefficient folded(3, [](const mfem::Vector &x, mfem::Vector &u) {
            u.SetSize(3); u = 0.0; u(0) = -2.0*x(0);
        });
        field.ProjectCoefficient(folded); field.GetTrueDofs(state);
        CHECK(compare() < 0.0);
        if (state.Size()) state(0) = std::numeric_limits<double>::quiet_NaN();
        CHECK_FALSE(std::isfinite(inspector.Inspect(space, state)));
        CHECK_FALSE(std::isfinite(inspector.Inspect(space, state, true)));
    }
}

TEST_CASE("Geometry Inspector Preserves Tetrahedron Fallback", tags::unit &tags::geometry) {
    auto serial = mfem::Mesh::MakeCartesian3D(1, 1, 1, mfem::Element::TETRAHEDRON);
    mfem::ParMesh mesh(MPI_COMM_WORLD, serial);
    mfem::H1_FECollection collection(1, 3);
    mfem::ParFiniteElementSpace space(&mesh, &collection, 3);
    mfem::Vector state(space.GetTrueVSize()); state = 0.0;
    mean_field::fem::PreparedGeometryInspector inspector;
    CHECK(inspector.Inspect(space, state) == 1.0);
    CHECK(inspector.Inspect(space, state, true) == 1.0);
    CHECK(inspector.GetTensorElementCount() == 0);
    CHECK(inspector.GetGeometryBuildCount() == 0);
}

TEST_CASE("Tensor Gradient Transposes Match Dense Assembly And Adjoint Products", tags::unit &tags::quadrature) {
    const auto check = [](const mfem::FiniteElement &element, const mfem::IntegrationRule &rule) {
        TableCache cache;
        const auto table = cache.GetScalarTable(element, rule);
        REQUIRE(table->HasTensorEvaluation3D());
        constexpr int components = 2;
        mfem::DenseMatrix dofs(element.GetDof(), components), dual(rule.GetNPoints(), 3*components),
                          values, actual, expected(element.GetDof(), components);
        mfem::Vector workspace;
        for (int c = 0; c < components; ++c) for (int i = 0; i < dofs.Height(); ++i)
            dofs(i,c) = std::sin(0.37*(i+1)+c);
        for (int c = 0; c < dual.Width(); ++c) for (int q = 0; q < dual.Height(); ++q)
            dual(q,c) = std::cos(0.23*(q+1)+0.31*c);

        for (int repeat = 0; repeat < 2; ++repeat) {
            if (repeat) dual = 0.0;
            table->EvaluateTensor3D(dofs, values, workspace);
            table->ApplyGradientTransposeTensor3D(dual, actual, workspace);
            expected = 0.0;
            double forwardProduct = 0.0, transposeProduct = 0.0, scale = 1.0;
            for (int c = 0; c < components; ++c) for (int q = 0; q < rule.GetNPoints(); ++q)
                for (int a = 0; a < 3; ++a) {
                    forwardProduct += values(q,4*c+a+1)*dual(q,3*c+a);
                    for (int i = 0; i < element.GetDof(); ++i)
                        expected(i,c) += table->GetGradients(q)(i,a)*dual(q,3*c+a);
                }
            for (int c = 0; c < components; ++c) for (int i = 0; i < dofs.Height(); ++i)
                scale = std::max(scale, std::abs(expected(i,c)));
            for (int c = 0; c < components; ++c) for (int i = 0; i < dofs.Height(); ++i) {
                CHECK(std::abs(actual(i,c)-expected(i,c)) <= 2.0e-11*scale);
                transposeProduct += dofs(i,c)*actual(i,c);
                if (repeat) CHECK(actual(i,c) == 0.0);
            }
            CHECK(std::abs(forwardProduct-transposeProduct) <=
                  2.0e-11*std::max(1.0, std::abs(forwardProduct)));
        }
    };
    for (int order : {1, 2, 4}) {
        CAPTURE(order);
        mfem::H1_FECollection collection(order, 3);
        check(*collection.FiniteElementForGeometry(mfem::Geometry::CUBE),
              mfem::IntRules.Get(mfem::Geometry::CUBE, 2*order+3));
    }
    for (int order : {0, 1, 3}) for (int basis : {mfem::BasisType::GaussLegendre, mfem::BasisType::OpenUniform}) {
        CAPTURE(order, basis);
        mfem::L2_FECollection collection(order, 3, basis);
        check(*collection.FiniteElementForGeometry(mfem::Geometry::CUBE),
              mfem::IntRules.Get(mfem::Geometry::CUBE, 2*order+3));
    }
}