#include <cmath>

#include <catch2/catch_test_macros.hpp>
#include "mfem/fem/fe_coll.hpp"
#include "mfem/mesh/mesh.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/backend/mfem/operators/operand_evaluation.hpp"
#include "serif/discretization/blocks/fields.hpp"

namespace {
    using Scalar = ::serif::tests::equations::Enthalpy;
    using Gradient = serif::discretization::forms::Operand<Scalar, serif::discretization::forms::Operation::Gradient>;
    using Evaluator = serif::operators::backend::mfem::OperandEvaluator<Gradient>;
    using serif::discretization::backend::mfem::BasisEvaluation;
    using serif::discretization::backend::mfem::BasisWorkspace;
    using serif::discretization::backend::mfem::ContractionMode;
    using serif::discretization::backend::mfem::TensorBasisCache;
}

TEST_CASE("Scalar H1 gradients agree with physical MFEM bases and transpose actions", "[equations][scalar_gradient][unit]") {
    for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
        auto mesh = ::mfem::Mesh::MakeCartesian3D(1, 1, 1, geometry, 1.7, 0.9, 1.2);
        // Shear as well as unequal lengths exposes inverse/transpose mistakes.
        for (int vertex = 0; vertex < mesh.GetNV(); ++vertex) {
            double* position = mesh.GetVertex(vertex);
            position[0] += 0.23 * position[1];
            position[2] -= 0.17 * position[1];
        }
        ::mfem::H1_FECollection collection(3, 3);
        TensorBasisCache cache;
        BasisWorkspace workspace;
        for (int index = 0; index < mesh.GetNE(); ++index) {
            auto& transformation = *mesh.GetElementTransformation(index);
            const auto& element = *collection.FiniteElementForGeometry(mesh.GetElementBaseGeometry(index));
            const auto& rule = ::mfem::IntRules.Get(element.GetGeomType(), 8);
            for (const auto mode : {ContractionMode::automatic, ContractionMode::dense}) {
                BasisEvaluation basis;
                Evaluator::sample_basis(element, rule, basis, cache, mode, transformation, workspace);
                ::mfem::Vector coefficients(element.GetDof()), values, dual(3 * rule.GetNPoints()), transpose;
                for (int dof = 0; dof < coefficients.Size(); ++dof) {
                    coefficients[dof] = std::sin(0.37 * (dof + 1));
                }
                Evaluator::evaluate(basis, coefficients, values, workspace);
                ::mfem::DenseMatrix physical_basis(element.GetDof(), 3);
                for (int point = 0; point < rule.GetNPoints(); ++point) {
                    transformation.SetIntPoint(&rule.IntPoint(point));
                    element.CalcPhysDShape(transformation, physical_basis);
                    for (int axis = 0; axis < 3; ++axis) {
                        double expected = 0.0;
                        for (int dof = 0; dof < coefficients.Size(); ++dof) {
                            expected += physical_basis(dof, axis) * coefficients[dof];
                        }
                        CHECK(std::abs(values[3 * point + axis] - expected) < 2.0e-11);
                        dual[3 * point + axis] = std::cos(0.19 * (3 * point + axis + 1));
                    }
                }
                Evaluator::integrate(basis, dual, transpose, workspace);
                CHECK(std::abs(values * dual - coefficients * transpose) < 2.0e-10);
            }
        }
    }
}

TEST_CASE("Scalar H1 gradient geometry derivatives follow inverse transpose mapping", "[equations][scalar_gradient][unit]") {
    serif::mapping::MappingPointContext context;
    context.mapping_jacobian.SetSize(3);
    context.mapping_jacobian = 0.0;
    context.mapping_jacobian(0, 0) = 1.2;
    context.mapping_jacobian(1, 1) = 0.8;
    context.mapping_jacobian(2, 2) = 1.1;
    context.mapping_jacobian(0, 1) = 0.21;
    context.mapping_jacobian(2, 0) = -0.13;
    context.inverse_mapping_jacobian.SetSize(3);
    ::mfem::CalcInverse(context.mapping_jacobian, context.inverse_mapping_jacobian);
    serif::mapping::MappingPointVariation variation;
    variation.mapping_jacobian_variation.SetSize(3);
    for (int row = 0; row < 3; ++row) {
        for (int column = 0; column < 3; ++column) {
            variation.mapping_jacobian_variation(row, column) = std::sin(0.41 * (3 * row + column + 1));
        }
    }
    ::mfem::DenseMatrix product(3);
    variation.inverse_mapping_jacobian_variation.SetSize(3);
    ::mfem::Mult(context.inverse_mapping_jacobian, variation.mapping_jacobian_variation, product);
    ::mfem::Mult(product, context.inverse_mapping_jacobian, variation.inverse_mapping_jacobian_variation);
    variation.inverse_mapping_jacobian_variation *= -1.0;
    const Evaluator::ReferenceValue reference{0.7, -0.2, 0.4};
    const auto physical = Evaluator::physical_value(reference, context);
    // A^T grad_x = grad_X is independent of the implementation's inverse multiplication.
    for (int axis = 0; axis < 3; ++axis) {
        double recovered = 0.0;
        for (int component = 0; component < 3; ++component) {
            recovered += context.mapping_jacobian(component, axis) * physical[component];
        }
        CHECK(std::abs(recovered - reference[axis]) < 2.0e-14);
    }
    const auto derivative = Evaluator::geometry_variation(reference, context, variation);
    auto plus = context, minus = context;
    constexpr double step = 1.0e-6;
    plus.mapping_jacobian.Add(step, variation.mapping_jacobian_variation);
    minus.mapping_jacobian.Add(-step, variation.mapping_jacobian_variation);
    ::mfem::CalcInverse(plus.mapping_jacobian, plus.inverse_mapping_jacobian);
    ::mfem::CalcInverse(minus.mapping_jacobian, minus.inverse_mapping_jacobian);
    const auto plus_gradient = Evaluator::physical_value(reference, plus);
    const auto minus_gradient = Evaluator::physical_value(reference, minus);
    for (int axis = 0; axis < 3; ++axis) {
        CHECK(std::abs(derivative[axis] - (plus_gradient[axis] - minus_gradient[axis]) / (2.0 * step)) < 2.0e-10);
    }
}