#include <cmath>

#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/fe_coll.hpp"
#include "serif/backend/mfem/discretization/basis_evaluation.hpp"

namespace
{
    using namespace serif::discretization::backend::mfem;

    /// Require repeatability after the original finite element, quadrature, and cache have been destroyed.
    void check_retained_basis(const int basis)
    {
        BasisEvaluation retained;
        BasisWorkspace workspace;
        ::mfem::Vector coefficients, dual, expected_values, expected_transpose;
        {
            const ::mfem::H1_FECollection collection(3, 3, basis);
            const ::mfem::FiniteElement& element = *collection.FiniteElementForGeometry(::mfem::Geometry::CUBE);
            const ::mfem::IntegrationRule rule(::mfem::IntRules.Get(::mfem::Geometry::CUBE, 7));
            TensorBasisCache cache;
            retained.sample(element, rule, ReferenceOperation::value_and_gradient, 1, cache, ContractionMode::automatic);
            CHECK(retained.uses_tensor() == (basis != ::mfem::BasisType::Positive));
            coefficients.SetSize(retained.width());
            dual.SetSize(retained.height());
            for (int i = 0; i < coefficients.Size(); ++i)
            {
                coefficients[i] = std::sin(0.37 * (i + 1));
            }
            for (int i = 0; i < dual.Size(); ++i)
            {
                dual[i] = std::cos(0.23 * (i + 1));
            }
            retained.evaluate(coefficients, expected_values, workspace);
            retained.integrate(dual, expected_transpose, workspace);
        }

        ::mfem::Vector values, transpose;
        retained.evaluate(coefficients, values, workspace);
        retained.integrate(dual, transpose, workspace);
        values -= expected_values;
        transpose -= expected_transpose;
        CHECK(values.Normlinf() == 0.0);
        CHECK(transpose.Normlinf() == 0.0);
    }
}

TEST_CASE("MeanField scalar basis actions survive cache finite element and rule destruction", "[equations][tensor_basis][unit]")
{
    check_retained_basis(::mfem::BasisType::GaussLobatto);
    check_retained_basis(::mfem::BasisType::Positive);
}

TEST_CASE("MeanField basis cache follows rule contents without mutating published factors", "[equations][tensor_basis][unit]")
{
    ::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 copy(rule), changed_weight(rule), moved_point(rule);
    changed_weight.IntPoint(0).weight *= 1.25;
    moved_point.IntPoint(0).x += 0.03125;
    TensorBasisCache cache;
    const auto original = cache.get(element, rule, ReferenceOperation::value_and_gradient, 1);
    REQUIRE(original);
    ::mfem::Vector coefficients(original->dofs()), expected, values;
    coefficients = 0.37;
    TensorWorkspace workspace;
    original->apply(coefficients, expected, workspace);

    CHECK(cache.get(element, copy, ReferenceOperation::value_and_gradient, 1) == original);
    const auto reweighted = cache.get(element, changed_weight, ReferenceOperation::value_and_gradient, 1);
    REQUIRE(reweighted);
    CHECK(reweighted != original);
    CHECK_FALSE(cache.get(element, moved_point, ReferenceOperation::value_and_gradient, 1));
    CHECK(cache.size() == 3);
    copy.IntPoint(0).x = moved_point.IntPoint(0).x;
    CHECK_FALSE(cache.get(element, copy, ReferenceOperation::value_and_gradient, 1));
    CHECK(cache.size() == 3);
    CHECK(cache.get(element, rule, ReferenceOperation::value_and_gradient, 1) == original);

    // Weights identify cache entries, but basis evaluation introduces no integration weight.
    reweighted->apply(coefficients, values, workspace);
    values -= expected;
    CHECK(values.Normlinf() == 0.0);
    cache.clear();
    CHECK(cache.size() == 0);
    original->apply(coefficients, values, workspace);
    values -= expected;
    CHECK(values.Normlinf() == 0.0);
}