#include <cmath>

#include <catch2/catch_test_macros.hpp>
#include "serif/tests/equations/fixture.hpp"

namespace {
    using namespace serif::tests::equations;
    using Scalar = ::serif::tests::equations::Enthalpy;
    using Gradient = serif::discretization::forms::Operand<Scalar, serif::discretization::forms::Operation::Gradient>;

    /** @brief Exercise a scalar diffusion row through the production volume assembly. */
    struct ScalarDiffusion : serif::equations::Contribution<Gradient> {
        using Operands = serif::utils::types::TypeList<Gradient>;
        static constexpr std::string_view name = "scalar gradient assembly regression";

        /** @brief The flux is grad(z), so its partial with respect to grad(z) is identity. */
        serif::equations::PointLinearization<Operands, serif::discretization::forms::VectorValue>
        evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            serif::equations::PointLinearization<Operands, serif::discretization::forms::VectorValue> result;
            result.residual = point.fields.get<Gradient>();
            for (int axis = 0; axis < 3; ++axis) {
                result.field_derivatives.get<Gradient>()[axis][axis] = 1.0;
            }
            return result;
        }

        /** @brief This linear flux introduces no extra polynomial coefficient. */
        static int additional_order(Gradient, int) { return 0; }
    };
}

TEST_CASE("Scalar gradient test rows assemble diffusion and differentiate moving geometry", "[equations][scalar_gradient][unit]") {
    for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
        Fixture fixture(geometry);
        const auto values = deterministic(fixture.enthalpy_binding.size(), 0.31);
        const auto direction = deterministic(values.Size(), 0.57);
        const auto displacement = fixture.displacement(0.6);
        const auto displacement_direction = fixture.displacement(0.2);
        const auto zero = constant(displacement.Size(), 0.0);
        const auto fields = [&](const ::mfem::Vector& coefficients) {
            return backend::make_term_coefficients(backend::bind(fixture.enthalpy_binding, coefficients));
        };
        serif::discretization::quadrature::RuleSet rules;
        rules.set_fallback_control({.fixed_order = 10});
        VolumeEquationFixture prepared{ScalarDiffusion{}, fixture.enthalpy_binding, fixture.displacement_binding,
            fixture.compactification, fixture.mapper, serif::discretization::quadrature::Policy{rules}};
        prepared.prepare(fields(values), zero, fixture.workspace);
        ::mfem::Vector residual;
        prepared.residual(fixture.workspace, residual);
        ::mfem::BilinearForm reference(&fixture.enthalpy_space);
        ::mfem::Array<int> stellar{1, 1, 0};
        reference.AddDomainIntegrator(new ::mfem::DiffusionIntegrator(), stellar);
        reference.Assemble();
        reference.Finalize();
        const auto full_values = fixture.enthalpy_binding.scatter(values);
        ::mfem::Vector full_residual(full_values.Size());
        reference.Mult(full_values, full_residual);
        CHECK(relative_error(residual, fixture.enthalpy_binding.gather(full_residual)) < 3.0e-12);

        prepared.prepare(fields(values), displacement, fixture.workspace);
        if (geometry == ::mfem::Element::HEXAHEDRON) {
            CHECK(prepared.tensor_element_count() > 0);
        }
        ::mfem::Vector action, transpose;
        prepared.apply_jacobian(fields(direction), displacement_direction, fixture.workspace, action);
        const auto dual = deterministic(residual.Size(), 0.83);
        prepared.apply_jacobian_transpose(dual, fixture.workspace, transpose);
        const auto packed_direction = packed(constant(0, 0.0), direction, displacement_direction);
        CHECK(std::abs(dual * action - packed_direction * transpose) < 3.0e-11);

        constexpr double step = 1.0e-6;
        auto plus = values, minus = values, plus_displacement = displacement, minus_displacement = displacement;
        plus.Add(step, direction);
        minus.Add(-step, direction);
        plus_displacement.Add(step, displacement_direction);
        minus_displacement.Add(-step, displacement_direction);
        ::mfem::Vector plus_residual, minus_residual;
        prepared.prepare(fields(plus), plus_displacement, fixture.workspace);
        prepared.residual(fixture.workspace, plus_residual);
        prepared.prepare(fields(minus), minus_displacement, fixture.workspace);
        prepared.residual(fixture.workspace, minus_residual);
        plus_residual -= minus_residual;
        plus_residual /= 2.0 * step;
        CHECK(relative_error(action, plus_residual) < 3.0e-8);
    }
}