#include <algorithm>
#include <array>
#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;

    /// Check every basis coefficient against MFEM, including bases that require dense fallback.
    void check_native_basis(const ::mfem::FiniteElement& element, const ::mfem::IntegrationRule& rule, TensorBasisCache& cache)
    {
        const bool vector = element.GetRangeType() == ::mfem::FiniteElement::VECTOR;
        const ReferenceOperation operation = vector ? ReferenceOperation::value : ReferenceOperation::value_and_gradient;
        const int components = vector ? 3 : 4;
        BasisEvaluation evaluation;
        evaluation.sample(element, rule, operation, 1, cache, ContractionMode::automatic);
        BasisWorkspace workspace;
        ::mfem::Vector coefficients(element.GetDof()), values, shape(element.GetDof());
        ::mfem::DenseMatrix derivatives(element.GetDof(), 3);

        for (int dof = 0; dof < element.GetDof(); ++dof)
        {
            coefficients = 0.0;
            coefficients[dof] = 1.0;
            evaluation.evaluate(coefficients, values, workspace);
            REQUIRE(std::all_of(values.GetData(), values.GetData() + values.Size(), [](const double value)
            {
                return std::isfinite(value);
            }));
            double error = 0.0;
            for (int point = 0; point < rule.GetNPoints(); ++point)
            {
                if (vector)
                {
                    element.CalcVShape(rule.IntPoint(point), derivatives);
                }
                else
                {
                    element.CalcShape(rule.IntPoint(point), shape);
                    element.CalcDShape(rule.IntPoint(point), derivatives);
                    error = std::max(error, std::abs(values[components * point] - shape[dof]));
                }

                for (int axis = 0; axis < 3; ++axis)
                {
                    const int sample = components * point + axis + (vector ? 0 : 1);
                    error = std::max(error, std::abs(values[sample] - derivatives(dof, axis)));
                }
            }
            CHECK(error < 2.0e-12);
        }
    }
}

TEST_CASE("MeanField native scalar reference values and gradients include positive and closed bases", "[equations][tensor_basis][unit]")
{
    for (const int order : {1, 3})
    {
        ::mfem::H1_FECollection h1(order, 3);
        ::mfem::L2_FECollection l2(order - 1, 3);
        TensorBasisCache cache;
        const ::mfem::IntegrationRule& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 2 * order + 1);
        check_native_basis(*h1.FiniteElementForGeometry(::mfem::Geometry::CUBE), rule, cache);
        check_native_basis(*l2.FiniteElementForGeometry(::mfem::Geometry::CUBE), rule, cache);
    }

    ::mfem::H1_FECollection nodal(3, 3), positive(3, 3, ::mfem::BasisType::Positive);
    ::mfem::L2_FECollection open(3, 3), closed(3, 3, ::mfem::BasisType::GaussLobatto);
    TensorBasisCache cache;
    const ::mfem::IntegrationRule& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 5);
    std::array<std::shared_ptr<const TensorBasis>, 4> plans;
    const std::array<const ::mfem::FiniteElementCollection*, 4> collections{&nodal, &positive, &open, &closed};
    for (std::size_t index = 0; index < collections.size(); ++index)
    {
        const ::mfem::FiniteElement& element = *collections[index]->FiniteElementForGeometry(::mfem::Geometry::CUBE);
        check_native_basis(element, rule, cache);
        plans[index] = cache.get(element, rule, ReferenceOperation::value_and_gradient, 1);
        if (plans[index])
        {
            for (std::size_t previous = 0; previous < index; ++previous)
            {
                CHECK(plans[index] != plans[previous]);
            }
        }
    }
}