#include <algorithm>
#include <cmath>
#include <catch2/catch_test_macros.hpp>
#include "mfem/fem/bilinearform.hpp"
#include "mfem/fem/bilininteg.hpp"
#include "serif/tests/equations/fixture.hpp"
#include "serif/tests/nonconforming.hpp"

namespace {
    using namespace serif::tests::equations;
    namespace forms = serif::discretization::forms;
    using Vector = ::mfem::Vector;
    using VectorField = blocks::displacement::vector;
    using Value = forms::Operand<VectorField>;
    using Gradient = forms::Operand<VectorField, forms::Operation::Gradient>;
    using StellarDomains = serif::discretization::domain::StellarDomains;

    /** @brief Supply the componentwise diffusion flux for comparison with native MFEM. */
    struct VectorDiffusion : serif::equations::Contribution<Gradient, StellarDomains> {
        using Operands = serif::utils::types::TypeList<Gradient>;
        static constexpr std::string_view name = "vector diffusion regression";
        /** @brief Each physical gradient component contributes its own identity derivative. */
        serif::equations::PointLinearization<Operands, forms::MatrixValue> evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            serif::equations::PointLinearization<Operands, forms::MatrixValue> result;
            result.residual = point.fields.get<Gradient>();
            for (int entry = 0; entry < 9; ++entry) {
                result.field_derivatives.get<Gradient>()[entry][entry] = 1.0;
            }
            return result;
        }
        /** @brief The diffusion flux has no additional polynomial coefficient. */
        static int additional_order(Gradient, int) { return 0; }
    };

    /** @brief Couple transposed gradient entries and vector values to expose tensor indexing errors. */
    struct CoupledFlux : serif::equations::Contribution<Gradient, StellarDomains> {
        using Operands = serif::utils::types::TypeList<Gradient, Value>;
        static constexpr std::string_view name = "coupled tensor flux regression";
        /** @brief Differentiate a nonlinear manufactured flux, including its explicit position dependence. */
        serif::equations::PointLinearization<Operands, forms::MatrixValue> evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            const auto& gradient = point.fields.get<Gradient>();
            const auto& value = point.fields.get<Value>();
            serif::equations::PointLinearization<Operands, forms::MatrixValue> result;
            for (int component = 0; component < 3; ++component) {
                for (int axis = 0; axis < 3; ++axis) {
                    const int entry = 3 * component + axis;
                    result.residual[component][axis] = gradient[component][axis] + 0.37 * gradient[axis][component] +
                        0.11 * gradient[component][axis] * gradient[component][axis] + 0.19 * value[component] * value[axis] +
                        0.07 * point.position[axis] * value[component];
                    auto& gradient_partial = result.field_derivatives.get<Gradient>();
                    gradient_partial[entry][entry] += 1.0 + 0.22 * gradient[component][axis];
                    gradient_partial[entry][3 * axis + component] += 0.37;
                    auto& value_partial = result.field_derivatives.get<Value>();
                    value_partial[entry][component] += 0.19 * value[axis] + 0.07 * point.position[axis];
                    value_partial[entry][axis] += 0.19 * value[component];
                    result.position_derivative[entry][axis] = 0.07 * value[component];
                }
            }
            return result;
        }
        /** @brief Account for the quadratic factors in this manufactured flux. */
        template <typename Operand> static int additional_order(Operand, int order) { return order + 1; }
    };

    /** @brief Exercise the reverse rectangular derivative from tensor input to vector output. */
    struct CoupledForce : serif::equations::Contribution<Value, StellarDomains> {
        using Operands = serif::utils::types::TypeList<Gradient, Value>;
        static constexpr std::string_view name = "coupled vector force regression";
        /** @brief Contract the gradient with velocity and differentiate both independent operands. */
        serif::equations::PointLinearization<Operands, forms::VectorValue> evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            const auto& gradient = point.fields.get<Gradient>();
            const auto& value = point.fields.get<Value>();
            serif::equations::PointLinearization<Operands, forms::VectorValue> result;
            for (int component = 0; component < 3; ++component) {
                for (int axis = 0; axis < 3; ++axis) {
                    result.residual[component] += gradient[component][axis] * value[axis];
                    result.field_derivatives.get<Gradient>()[component][3 * component + axis] = value[axis];
                    result.field_derivatives.get<Value>()[component][axis] = gradient[component][axis];
                }
            }
            return result;
        }
        /** @brief The convection-like product adds one field's polynomial order. */
        template <typename Operand> static int additional_order(Operand, int order) { return order; }
    };

    /** @brief Publish the flux and force as two contributions of one vector row. */
    struct CoupledEquation {
        using Test = VectorField;
        static constexpr std::string_view name = "coupled tensor flux and vector force regression";
        [[nodiscard]] auto contributions() const { return std::tuple{CoupledFlux{}, CoupledForce{}}; }
    };

    /** @brief Fix identical quadrature for native MFEM and the prepared equation. */
    serif::discretization::quadrature::Policy integration_policy() {
        serif::discretization::quadrature::RuleSet rules;
        rules.set_fallback_control({.fixed_order = 12});
        return serif::discretization::quadrature::Policy{rules};
    }

    /** @brief Give the physical vector field a coefficient input independent of mesh displacement. */
    auto field(const Fixture& fixture, const Vector& values) {
        return backend::make_term_coefficients(backend::bind(fixture.displacement_binding, values));
    }

    /** @brief Compare native local assembly followed by P transpose with the reduced backend residual. */
    void check_native_diffusion(Fixture& fixture, backend::ContractionMode mode) {
        const Vector values = deterministic(fixture.displacement_binding.size(), 0.31);
        const Vector zero = constant(values.Size(), 0.0);
        VolumeEquationFixture prepared{VectorDiffusion{}, fixture.displacement_binding, fixture.displacement_binding,
            fixture.compactification, fixture.mapper, integration_policy(), mode};
        prepared.prepare(field(fixture, values), zero, fixture.workspace);
        Vector actual;
        prepared.residual(fixture.workspace, actual);
        ::mfem::BilinearForm reference(&fixture.displacement_space);
        auto* integrator = new ::mfem::VectorDiffusionIntegrator();
        integrator->SetIntRule(&::mfem::IntRules.Get(fixture.mesh.GetElementBaseGeometry(0), 12));
        ::mfem::Array<int> stellar{1, 1, 0};
        reference.AddDomainIntegrator(integrator, stellar);
        reference.Assemble();
        reference.Finalize();
        const Vector full_values = fixture.displacement_binding.scatter(values);
        Vector native(full_values.Size());
        reference.Mult(full_values, native);
        CHECK(relative_error(actual, fixture.displacement_binding.accumulate_dual(native)) < 2.0e-11);
        if (fixture.mesh.GetElementBaseGeometry(0) == ::mfem::Geometry::CUBE) {
            CHECK((prepared.tensor_element_count() > 0) == (mode == backend::ContractionMode::automatic));
        }
    }

    /** @brief Check each independent field/geometry column, combined action, and complete transpose. */
    void check_coupled_derivatives(Fixture& fixture, backend::ContractionMode mode) {
        const Vector values = deterministic(fixture.displacement_binding.size(), 0.31);
        const Vector direction = deterministic(values.Size(), 0.57);
        const Vector displacement = fixture.displacement(0.6);
        const Vector displacement_direction = fixture.displacement(0.2);
        const CoupledEquation equation;
        VolumeEquationFixture prepared{equation, fixture.displacement_binding, fixture.displacement_binding,
            fixture.compactification, fixture.mapper, integration_policy(), mode};
        const Vector dual = deterministic(fixture.displacement_binding.size(), 0.83);
        for (int column = 0; column < 3; ++column) {
            Vector selected = direction, selected_geometry = displacement_direction;
            if (column == 0) { selected_geometry = 0.0; }
            if (column == 1) { selected = 0.0; }
            prepared.prepare(field(fixture, values), displacement, fixture.workspace);
            Vector action, transpose;
            prepared.apply_jacobian(field(fixture, selected), selected_geometry, fixture.workspace, action);
            prepared.apply_jacobian_transpose(dual, fixture.workspace, transpose);
            const Vector packed_direction = packed(constant(0, 0.0), selected, selected_geometry);
            const double forward = dual * action;
            CHECK(std::abs(forward - packed_direction * transpose) < 3.0e-11 * std::max(1.0, std::abs(forward)));
            // Geometry directions are small; these steps avoid subtractive cancellation without relaxing the tolerance.
            for (const double step : {1.0e-4, 1.0e-5}) {
                CAPTURE(column, step);
                Vector plus = values, minus = values, plus_geometry = displacement, minus_geometry = displacement;
                plus.Add(step, selected);
                minus.Add(-step, selected);
                plus_geometry.Add(step, selected_geometry);
                minus_geometry.Add(-step, selected_geometry);
                Vector plus_residual, minus_residual;
                prepared.prepare(field(fixture, plus), plus_geometry, fixture.workspace);
                prepared.residual(fixture.workspace, plus_residual);
                prepared.prepare(field(fixture, minus), minus_geometry, fixture.workspace);
                prepared.residual(fixture.workspace, minus_residual);
                plus_residual -= minus_residual;
                plus_residual /= 2.0 * step;
                CHECK(relative_error(action, plus_residual) < 5.0e-8);
            }
        }
    }
}

TEST_CASE("Vector gradient equations retain dense tensor and ordering equivalence", "[equations][vector_gradient][unit]") {
    for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
        for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
            Fixture fixture(geometry, ordering);
            for (const auto mode : {backend::ContractionMode::automatic, backend::ContractionMode::dense}) {
                CAPTURE(geometry, ordering, mode);
                check_native_diffusion(fixture, mode);
                check_coupled_derivatives(fixture, mode);
            }
        }
    }
}

TEST_CASE("Vector gradient equations preserve true-coordinate work on STROID hanging meshes", "[equations][vector_gradient][integration][nonconforming]") {
    for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
        auto generated = serif::tests::nonconforming::mesh(0, 3);
        REQUIRE(generated.mesh->Nonconforming());
        Fixture fixture(::mfem::Mesh(*generated.mesh), ordering);
        REQUIRE(fixture.displacement_space.GetTrueVSize() < fixture.displacement_space.GetVSize());
        CAPTURE(ordering);
        check_native_diffusion(fixture, backend::ContractionMode::automatic);
        check_coupled_derivatives(fixture, backend::ContractionMode::automatic);
    }
}