#include <algorithm>
#include <limits>

#include "serif/tests/equations/gravity_port.hpp"

namespace {
    using namespace serif;
    using namespace serif::tests::equations;
    using gravity::GravityFixture;
    using gravity_port::project_vector;
    using Vector = ::mfem::Vector;
    using Displacement = blocks::displacement::vector;
    using Divergence = discretization::forms::Operand<Displacement, discretization::forms::Operation::Divergence>;

    struct DivergenceInput {
        using Test = blocks::density::mass;
        using Operands = utils::types::TypeList<Divergence>;
        static constexpr std::string_view name = "H1 divergence input test";

        serif::equations::PointLinearization<Operands> evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            serif::equations::PointLinearization<Operands> result;
            result.residual = point.fields.get<Divergence>();
            result.field_derivatives.get<Divergence>() = 1.0;

            return result;
        }

        int additional_order(Divergence, int) const {
            return 0;
        }
    };

    /// A smooth nonlinear displacement direction, independent of the field being differentiated.
    Vector direction_displacement(GravityFixture& fixture) {
        return project_vector(fixture.displacement_binding, [](const Vector& position, Vector& value) {
            const double x = position[0], y = position[1], z = position[2];
            value.SetSize(3);
            value[0] = 0.019 * x + 0.011 * y * z - 0.006 * z * z + 0.004 * 0.83 * y;
            value[1] = -0.016 * y + 0.008 * x * z + 0.005 * x * x - 0.003 * 0.83 * z;
            value[2] = 0.013 * z - 0.010 * x * y + 0.006 * y * y + 0.004 * 0.83 * x;
        });
    }

    double relative(const Vector& actual, const Vector& expected) {
        Vector error(actual);
        error -= expected;

        return error.Norml2() / std::max({actual.Norml2(), expected.Norml2(), 100.0 * std::numeric_limits<double>::epsilon()});
    }
}

TEST_CASE("H1 divergence inputs retain component gradients under a nonsymmetric map", "[equations][h1_divergence][unit]") {
    for (const auto geometry : {::mfem::Element::TETRAHEDRON, ::mfem::Element::HEXAHEDRON}) {
        for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
            GravityFixture fixture(geometry, ordering);
            const auto value = project_vector(fixture.displacement_binding, [](const Vector& x, Vector& u) {
                u.SetSize(3);
                u[0] = 0.4 * x[0] - 0.2 * x[1] + 0.3 * x[2];
                u[1] = 0.1 * x[0] + 0.5 * x[1] - 0.1 * x[2];
                u[2] = -0.3 * x[0] + 0.2 * x[1] + 0.6 * x[2];
            });
            const auto geometry_value = project_vector(fixture.displacement_binding, [](const Vector& x, Vector& u) {
                u.SetSize(3);
                u[0] = 0.2 * x[0] + 0.07 * x[1];
                u[1] = -0.04 * x[0] - 0.1 * x[1] + 0.05 * x[2];
                u[2] = 0.03 * x[0] + 0.1 * x[2];
            });
            const auto direction = gravity_port::prepared_displacement(fixture, 0.83);
            const auto geometry_direction = direction_displacement(fixture);
            auto state = [&](const Vector& v) {
                return backend::make_term_coefficients(backend::bind(fixture.displacement_binding, v));
            };
            VolumeEquationFixture prepared{
                DivergenceInput{}, fixture.density_binding, fixture.displacement_binding,
                fixture.compactification, fixture.mapper, gravity::reference_policy()
            };
            prepared.prepare(state(value), geometry_value, fixture.workspace);
            Vector actual;
            prepared.residual(fixture.workspace, actual);

            ::mfem::DenseMatrix A(3), inverse(3);
            A = 0.0;
            A(0, 0) = 1.2;
            A(0, 1) = 0.07;
            A(1, 0) = -0.04;
            A(1, 1) = 0.9;
            A(1, 2) = 0.05;
            A(2, 0) = 0.03;
            A(2, 2) = 1.1;
            ::mfem::CalcInverse(A, inverse);
            constexpr double gradient[3][3] = {{0.4, -0.2, 0.3}, {0.1, 0.5, -0.1}, {-0.3, 0.2, 0.6}};
            double divergence = 0.0;
            for (int component = 0; component < 3; ++component) {
                for (int derivative = 0; derivative < 3; ++derivative) {
                    divergence += gradient[component][derivative] * inverse(derivative, component);
                }
            }

            ::mfem::ConstantCoefficient coefficient(divergence * A.Det());
            ::mfem::LinearForm reference_load(&fixture.density_space);
            ::mfem::Array<int> stellar{1, 1, 0};
            reference_load.AddDomainIntegrator(new ::mfem::DomainLFIntegrator(coefficient), stellar);
            reference_load.Assemble();
            CHECK(relative(actual, fixture.density_binding.gather(reference_load)) < 2.0e-12);

            for (int column = 0; column < 3; ++column) {
                auto selected = direction, selected_geometry = geometry_direction;
                if (column == 0) {
                    selected_geometry = 0.0;
                } else if (column == 1) {
                    selected = 0.0;
                }

                prepared.prepare(state(value), geometry_value, fixture.workspace);
                Vector derivative;
                prepared.apply_jacobian(state(selected), selected_geometry, fixture.workspace, derivative);
                auto plus = value, minus = value, plus_geometry = geometry_value, minus_geometry = geometry_value;
                constexpr double step = 1.0e-5;
                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(state(plus), plus_geometry, fixture.workspace);
                prepared.residual(fixture.workspace, plus_residual);
                prepared.prepare(state(minus), minus_geometry, fixture.workspace);
                prepared.residual(fixture.workspace, minus_residual);
                plus_residual -= minus_residual;
                plus_residual *= 0.5 / step;
                CHECK(relative(derivative, plus_residual) < 3.0e-8);
            }
        }
    }
}