#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;

    /// Compare prepared tensor or mapped actions against independently sampled dense matrices.
    void check_actions(const BasisEvaluation& tensor, const BasisEvaluation& dense, const double tolerance) {
        ::mfem::Vector dofs(tensor.width()), dual(tensor.height()), values, reference, assembled, reference_transpose;
        for (int i = 0; i < dofs.Size(); ++i) {
            dofs[i] = std::sin(0.53 * (i + 1));
        }
        for (int i = 0; i < dual.Size(); ++i) {
            dual[i] = std::cos(0.31 * (i + 1));
        }
        BasisWorkspace workspace;
        for (int repeat = 0; repeat < 2; ++repeat) {
            if (repeat) {
                dofs = 0.0;
                dual = 0.0;
            }
            tensor.evaluate(dofs, values, workspace);
            dense.evaluate(dofs, reference, workspace);
            ::mfem::Vector error = values;
            error -= reference;
            CHECK(error.Normlinf() <= tolerance * std::max(1.0, reference.Normlinf()));
            tensor.integrate(dual, assembled, workspace);
            dense.integrate(dual, reference_transpose, workspace);
            error = assembled;
            error -= reference_transpose;
            CHECK(error.Norml2() <= tolerance * std::max(1.0, reference_transpose.Norml2()));
            CHECK(std::abs(values * dual - dofs * assembled) <= tolerance * std::max(1.0, std::abs(values * dual)));
            if (repeat) {
                CHECK(values.Norml2() == 0.0);
                CHECK(assembled.Norml2() == 0.0);
            }
        }
    }

    /// Compare independent native MFEM basis matrices with the factored contractions and their transposes.
    void check_basis(const ::mfem::FiniteElement& element, const ::mfem::IntegrationRule& rule, ReferenceOperation operation, int copies, double tolerance) {
        TensorBasisCache cache;
        BasisEvaluation tensor, dense;
        tensor.sample(element, rule, operation, copies, cache, ContractionMode::automatic);
        dense.sample(element, rule, operation, copies, cache, ContractionMode::dense);
        REQUIRE(tensor.uses_tensor());
        REQUIRE_FALSE(dense.uses_tensor());
        check_actions(tensor, dense, tolerance);
    }

}

TEST_CASE("Scalar and vector H1 tensor values gradients and transposes retain MeanField accuracy", "[equations][tensor_basis][unit]")
{
    // MeanField tests/fem/reference_tables.cpp: retain its orders, bases and thresholds.
    for (const int order : {1, 2, 4})
    {
        for (const int basis : {::mfem::BasisType::GaussLobatto, ::mfem::BasisType::ClosedUniform})
        {
            ::mfem::H1_FECollection h1(order, 3, basis);
            const auto& element = *h1.FiniteElementForGeometry(::mfem::Geometry::CUBE);
            const auto& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 2 * order + 3);
            for (const int copies : {1, 2, 3})
            {
                check_basis(element, rule, ReferenceOperation::value, copies, 2.0e-12);
                check_basis(element, rule, ReferenceOperation::gradient, copies, 2.0e-11);
                check_basis(element, rule, ReferenceOperation::value_and_gradient, copies, 2.0e-12);
            }
        }
    }
    for (const int order : {0, 1, 3})
    {
        for (const int basis : {::mfem::BasisType::GaussLegendre, ::mfem::BasisType::OpenUniform})
        {
            ::mfem::L2_FECollection l2(order, 3, basis);
            const auto& element = *l2.FiniteElementForGeometry(::mfem::Geometry::CUBE);
            const auto& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 2 * order + 3);
            for (const int copies : {1, 2, 3})
            {
                check_basis(element, rule, ReferenceOperation::value, copies, 2.0e-12);
                check_basis(element, rule, ReferenceOperation::gradient, copies, 2.0e-11);
            }
        }
    }
}

TEST_CASE("Tensor plans own their factors and reject non Cartesian rules without changing quadrature", "[equations][tensor_basis][unit]")
{
    ::mfem::H1_FECollection h1(3, 3);
    const auto& element = *h1.FiniteElementForGeometry(::mfem::Geometry::CUBE);
    const auto& original = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 9);
    ::mfem::IntegrationRule copy(original.GetNPoints());
    for (int q = 0; q < copy.GetNPoints(); ++q)
    {
        copy.IntPoint(q) = original.IntPoint(q);
    }
    TensorBasisCache cache;
    auto retained = cache.get(element, copy, ReferenceOperation::value, 1);
    REQUIRE(retained);
    CHECK(cache.get(element, original, ReferenceOperation::value, 1) == retained);
    CHECK(cache.size() == 1);
    copy.IntPoint(copy.GetNPoints() - 1).x += 0.017;
    CHECK_FALSE(cache.get(element, copy, ReferenceOperation::value, 1));
    CHECK(cache.size() == 2);
    cache.clear();
    CHECK(cache.size() == 0);
    ::mfem::Vector dofs(retained->dofs()), values;
    dofs = 1.0;
    TensorWorkspace workspace;
    retained->apply(dofs, values, workspace);
    for (int q = 0; q < values.Size(); ++q)
    {
        CHECK(std::abs(values[q] - 1.0) < 2.0e-12);
    }

    const auto& tetrahedron = *h1.FiniteElementForGeometry(::mfem::Geometry::TETRAHEDRON);
    const auto& tetrahedron_rule = ::mfem::IntRules.Get(::mfem::Geometry::TETRAHEDRON, 9);
    CHECK_FALSE(TensorBasis::make(tetrahedron, tetrahedron_rule, ReferenceOperation::value, 1));
    CHECK_FALSE(TensorBasis::make(element, original, ReferenceOperation::value, 0));
    BasisEvaluation fallback;
    fallback.sample(element, copy, ReferenceOperation::value, 1, cache, ContractionMode::automatic);
    CHECK_FALSE(fallback.uses_tensor());
}

TEST_CASE("Reused basis actions retain mappings through composition resampling and dense fallback", "[equations][tensor_basis][unit]") {
    ::mfem::H1_FECollection h1(3, 3);
    const auto& element = *h1.FiniteElementForGeometry(::mfem::Geometry::CUBE);
    const auto& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 9);
    TensorBasisCache cache;
    BasisEvaluation reference, dense_reference;
    reference.sample(element, rule, ReferenceOperation::value_and_gradient, 3, cache, ContractionMode::automatic);
    dense_reference.sample(element, rule, ReferenceOperation::value_and_gradient, 3, cache, ContractionMode::dense);
    BasisEvaluation tensor = reference;
    BasisEvaluation dense = dense_reference;
    BasisWorkspace workspace;

    // Direct mapping of a sampled basis must disable its canonical shortcut, even without reset_from.
    for (int point = 0; point < rule.GetNPoints(); ++point) {
        const auto map = [point](const ::mfem::Vector& input, ::mfem::Vector& output) {
            for (int component = 0; component < output.Size(); ++component) {
                output[component] = (1.0 + 0.01 * point) * input[component] - 0.2 * input[(component + 1) % input.Size()];
            }
        };
        tensor.map_point(point, reference, map, workspace);
        dense.map_point(point, dense_reference, map, workspace);
    }
    check_actions(tensor, dense, 2.0e-12);

    // A mapped reference remains usable when the next physical map changes component count.
    reference = tensor;
    dense_reference = dense;
    tensor.reset_from(reference, 5);
    dense.reset_from(dense_reference, 5);
    for (int point = 0; point < rule.GetNPoints(); ++point) {
        const auto map = [](const ::mfem::Vector& input, ::mfem::Vector& output) {
            for (int component = 0; component < output.Size(); ++component) {
                output[component] = 0.7 * input[component] + 0.4 * input[component + 5];
            }
        };
        tensor.map_point(point, reference, map, workspace);
        dense.map_point(point, dense_reference, map, workspace);
    }
    check_actions(tensor, dense, 2.0e-12);

    // Reusing the same object for canonical, dense and constant bases must clear the earlier representation.
    for (const ContractionMode mode : {ContractionMode::automatic, ContractionMode::dense, ContractionMode::automatic}) {
        tensor.sample(element, rule, ReferenceOperation::gradient, 2, cache, mode);
        dense.sample(element, rule, ReferenceOperation::gradient, 2, cache, ContractionMode::dense);
        check_actions(tensor, dense, 2.0e-11);
    }
    tensor.set_constant(rule.GetNPoints());
    dense.set_constant(rule.GetNPoints());
    check_actions(tensor, dense, 2.0e-12);
}

TEST_CASE("Canonical mapped and dense basis actions preserve aliased inputs while resizing", "[equations][tensor_basis][unit]") {
    ::mfem::H1_FECollection h1(2, 3);
    const auto& element = *h1.FiniteElementForGeometry(::mfem::Geometry::CUBE);
    TensorBasisCache cache;
    BasisWorkspace workspace;
    for (const int order : {5, 7}) {
        const auto& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, order);
        for (const auto mode : {ContractionMode::automatic, ContractionMode::dense}) {
            CAPTURE(order, mode);
            BasisEvaluation reference, dense_reference;
            reference.sample(element, rule, ReferenceOperation::value, 1, cache, mode);
            dense_reference.sample(element, rule, ReferenceOperation::value, 1, cache, ContractionMode::dense);
            ::mfem::Vector input(reference.width()), dual(reference.height()), expected, aliased;
            for (int index = 0; index < input.Size(); ++index) {
                input[index] = std::sin(0.53 * (index + 1));
            }
            for (int index = 0; index < dual.Size(); ++index) {
                dual[index] = std::cos(0.31 * (index + 1));
            }
            for (const bool mapped : {false, true}) {
                BasisEvaluation basis = reference;
                BasisEvaluation dense = dense_reference;
                if (mapped) {
                    for (int point = 0; point < rule.GetNPoints(); ++point) {
                        const auto map = [point](const ::mfem::Vector& source, ::mfem::Vector& destination) {
                            destination[0] = (1.0 + 0.03 * point) * source[0];
                        };
                        basis.map_point(point, reference, map, workspace);
                        dense.map_point(point, dense_reference, map, workspace);
                    }
                }
                dense.evaluate(input, expected, workspace);
                aliased = input;
                basis.evaluate(aliased, aliased, workspace);
                aliased -= expected;
                CHECK(aliased.Normlinf() <= 2.0e-12);
                dense.integrate(dual, expected, workspace);
                aliased = dual;
                basis.integrate(aliased, aliased, workspace);
                aliased -= expected;
                CHECK(aliased.Normlinf() <= 2.0e-12);
            }
        }
    }
}

TEST_CASE("Implicit canonical identities preserve untouched points and support in place maps", "[equations][tensor_basis][unit]") {
    ::mfem::H1_FECollection h1(2, 3);
    const auto& element = *h1.FiniteElementForGeometry(::mfem::Geometry::CUBE);
    const auto& rule = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 7);
    TensorBasisCache cache;
    BasisWorkspace workspace;
    BasisEvaluation tensor, dense;
    // More than twelve components ensures mapping scratch does not restrict generality.
    tensor.sample(element, rule, ReferenceOperation::value_and_gradient, 4, cache, ContractionMode::automatic);
    dense.sample(element, rule, ReferenceOperation::value_and_gradient, 4, cache, ContractionMode::dense);
    REQUIRE(tensor.uses_tensor());
    const auto map = [](const ::mfem::Vector& input, ::mfem::Vector& output) {
        for (int component = 0; component < output.Size(); ++component) {
            output[component] = 0.8 * input[component] + 0.3 * input[(component + 1) % input.Size()];
        }
    };
    // A first map materializes only the destination representation; later maps compose with it.
    for (const int point : {3, 0, 3}) {
        tensor.map_point(point, tensor, map, workspace);
        dense.map_point(point, dense, map, workspace);
        check_actions(tensor, dense, 2.0e-11);
    }
}