// Ported from MeanField tests/deformation/safe_newton_step.cpp, with serial input and sampling regressions.

#include <array>
#include <cmath>
#include <limits>
#include <memory>
#include <utility>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include "mfem/fem/coefficient.hpp"
#include "mfem/fem/fe_coll.hpp"
#include "mfem/fem/gridfunc.hpp"
#include "mfem/mesh/mesh.hpp"
#include "serif/backend/mfem/geometry/deformation/geometry_step.hpp"
#include "serif/geometry/deformation/detail/determinant_polynomial.hpp"
#include "serif/backend/mfem/mapping/domain_mapper.hpp"

namespace {
    namespace deformation = serif::geometry::deformation;
    namespace backend = deformation::backend::mfem;
    namespace mapping = serif::mapping;
    using Catch::Matchers::WithinAbs;

    constexpr int dimension = 3;

    /** @brief Keep unit element widths so the mapped determinant has an independent analytic value. */
    [[nodiscard]] mfem::Mesh make_mesh(const bool compactified, const bool nonconforming = false) {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON, 2.0, 1.0, 1.0);
        if (nonconforming) {
            mesh.EnsureNCMesh();
            mfem::Array<int> selected{0};
            mesh.GeneralRefinement(selected);
            mesh.SetCurvature(6);
            mesh.Transform([](const mfem::Vector& x, mfem::Vector& y) {
                y = x;
                y[0] += 0.03 * x[0] * x[1];
                y[1] += 0.02 * x[1] * x[2];
                y[2] += 0.01 * x[0] * x[2];
            });
        }
        for (int element = 0; element < mesh.GetNE(); ++element) {
            mesh.GetElement(element)->SetAttribute(compactified ? 2 : 1);
        }
        return mesh;
    }

    /** @brief Reuse the physical Kelvin implementation in both interior and vacuum tests. */
    [[nodiscard]] std::unique_ptr<const mapping::compactification::ExteriorDomainMap> make_exterior_map() {
        return std::make_unique<mapping::compactification::KelvinCompactification>(serif::utils::user_options::KelvinCompactificationOptions{.r_star_ref = 1.0, .r_inf_ref = 4.0});
    }

    struct GeometryFixture {
        mfem::Mesh mesh;
        mfem::H1_FECollection displacement_collection;
        mfem::FiniteElementSpace displacement_space;
        mfem::H1_FECollection compactification_collection;
        mfem::FiniteElementSpace compactification_space;
        mfem::GridFunction compactification_coordinate;
        mapping::DomainMapper mapper;

        /** @brief Vary global ordering and interpolation degree without changing the reference geometry. */
        explicit GeometryFixture(const bool compactified = false, const mfem::Ordering::Type ordering = mfem::Ordering::byNODES, const int degree = 1, const bool nonconforming = false) :
            mesh(make_mesh(compactified, nonconforming)),
            displacement_collection(degree, dimension),
            displacement_space(&mesh, &displacement_collection, dimension, ordering),
            compactification_collection(1, dimension),
            compactification_space(&mesh, &compactification_collection),
            compactification_coordinate(&compactification_space),
            mapper({.dimension = dimension, .vacuum_element_attribute_ID = 2}, make_exterior_map()) {
            compactification_coordinate = 0.0;
        }

        /** @brief Provide an accepted identity map in true-DOF ordering. */
        [[nodiscard]] mfem::Vector zero_displacement() const {
            mfem::Vector result(displacement_space.GetTrueVSize());
            result = 0.0;
            return result;
        }

        /** @brief Construct directions independently of the geometry-step element extraction. */
        template <typename Function>
        [[nodiscard]] mfem::Vector project_direction(Function&& function) {
            mfem::VectorFunctionCoefficient coefficient(dimension, std::forward<Function>(function));
            mfem::GridFunction field(&displacement_space);
            field.ProjectCoefficient(coefficient);
            mfem::Vector result;
            field.GetTrueDofs(result);
            return result;
        }

        /** @brief Supply explicit sample rules for every element instead of choosing a hidden quadrature policy. */
        [[nodiscard]] std::vector<backend::GeometrySampleRule> sample_rules() {
            std::vector<backend::GeometrySampleRule> result;
            for (int element = 0; element < mesh.GetNE(); ++element) {
                result.push_back({.element = element, .integration_rule = &mfem::IntRules.Get(mesh.GetElementBaseGeometry(element), 2)});
            }
            return result;
        }
    };

    /** @brief Give det(J(alpha)) = 1 - 2 alpha. */
    void compress_x(const mfem::Vector& position, mfem::Vector& value) {
        value.SetSize(dimension);
        value = 0.0;
        value(0) = -2.0 * position(0);
    }

    /** @brief Give a tangent singularity at one half while both endpoints remain positive. */
    void compress_x_and_y(const mfem::Vector& position, mfem::Vector& value) {
        compress_x(position, value);
        value(1) = -2.0 * position(1);
    }

    /** @brief Provide a path with no geometry obstruction at positive steps. */
    void expand_x(const mfem::Vector& position, mfem::Vector& value) {
        value.SetSize(dimension);
        value = 0.0;
        value(0) = position(0);
    }

    class NonAffineExteriorMap final : public mapping::compactification::ExteriorDomainMap {
    public:
        /** @brief Declare that a cubic determinant model cannot certify this map. */
        [[nodiscard]] bool is_affine_in_displacement() const noexcept override {
            return false;
        }

        /** @brief Fail loudly if validation mistakenly reaches exterior evaluation. */
        [[nodiscard]] mapping::MappingStatus evaluate(const mapping::compactification::ExteriorMapInput&, mapping::compactification::ExteriorMapResult&) const override {
            FAIL("A non-affine exterior map must be rejected before evaluation.");
            return mapping::MappingStatus::NON_FINITE_MAPPING_RESULT;
        }

        /** @brief Fail loudly if validation mistakenly reaches exterior differentiation. */
        [[nodiscard]] mapping::MappingStatus evaluate_variation(const mapping::compactification::ExteriorMapInput&, const mapping::compactification::ExteriorMapResult&, const mapping::compactification::ExteriorMapDirection&, mapping::compactification::ExteriorMapVariation&) const override {
            FAIL("A non-affine exterior map must be rejected before differentiation.");
            return mapping::MappingStatus::NON_FINITE_MAPPING_RESULT;
        }

        /** @brief Identify the deliberately unsupported model in diagnostics. */
        [[nodiscard]] std::string_view get_name() const noexcept override {
            return "NonAffineExteriorMap";
        }
    };
}

TEST_CASE("Geometry step finds the first mapping boundary for both global DOF orderings", "[unit][deformation][geometry_step]") {
    for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM}) {
        GeometryFixture fixture(false, ordering);
        const mfem::Vector accepted = fixture.zero_displacement();
        const mfem::Vector direction = fixture.project_direction(compress_x);
        const auto rules = fixture.sample_rules();
        const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, accepted, direction, rules, {.fraction_to_boundary = 0.8});

        REQUIRE(result);
        CHECK(result->limited_by_geometry);
        CHECK_THAT(result->boundary_step, WithinAbs(0.5, 2.0e-13));
        CHECK_THAT(result->step, WithinAbs(0.4, 2.0e-13));
        CHECK_THAT(result->minimum_determinant_at_accepted, WithinAbs(1.0, 2.0e-13));
        CHECK_THAT(result->minimum_determinant_at_maximum, WithinAbs(-1.0, 2.0e-13));
        CHECK_THAT(result->limiting_point_determinant_at_step, WithinAbs(0.2, 2.0e-13));
        CHECK(result->sampled_point_count > 0);
        CHECK(result->limiting_element >= 0);
        CHECK(result->limiting_rule >= 0);
        CHECK(result->limiting_point >= 0);
    }
}

TEST_CASE("Geometry step detects a tangent singularity before an admissible endpoint", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture(true);
    const mfem::Vector accepted = fixture.zero_displacement();
    const mfem::Vector direction = fixture.project_direction(compress_x_and_y);
    const auto rules = fixture.sample_rules();
    const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, accepted, direction, rules);

    REQUIRE(result);
    CHECK(result->limited_by_geometry);
    CHECK_THAT(result->minimum_determinant_at_maximum, WithinAbs(1.0, 3.0e-13));
    CHECK_THAT(result->boundary_step, WithinAbs(0.5, 3.0e-13));
    CHECK_THAT(result->step, WithinAbs(0.45, 3.0e-13));
    CHECK_THAT(result->limiting_point_determinant_at_step, WithinAbs(0.01, 3.0e-13));
}

TEST_CASE("Geometry step preserves constrained displacement and all samples on curved NC geometry", "[unit][deformation][geometry_step][geometry_reuse]") {
    for (const bool compactified : {false, true}) {
        for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM}) {
            GeometryFixture fixture(compactified, ordering, 3, true);
            REQUIRE(fixture.displacement_space.GetTrueVSize() < fixture.displacement_space.GetVSize());
            const auto accepted = fixture.zero_displacement();
            const auto direction = fixture.project_direction(compress_x);
            auto rules = fixture.sample_rules();
            std::size_t expected_samples = 0;
            for (int element = 0; element < fixture.mesh.GetNE(); ++element) {
                rules.push_back({.element = element, .integration_rule = &mfem::IntRules.Get(fixture.mesh.GetElementBaseGeometry(element), 4)});
            }
            for (const auto& rule : rules) { expected_samples += rule.integration_rule->GetNPoints(); }
            const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space,
                fixture.compactification_coordinate, accepted, direction, rules, {.fraction_to_boundary = 0.8});
            REQUIRE(result);
            CHECK(result->sampled_point_count == expected_samples);
            CHECK(result->limited_by_geometry);
            CHECK_THAT(result->boundary_step, WithinAbs(0.5, 3.0e-12));
            CHECK_THAT(result->step, WithinAbs(0.4, 3.0e-12));
            CHECK_THAT(result->minimum_determinant_at_accepted, WithinAbs(1.0, 3.0e-12));
            CHECK_THAT(result->minimum_determinant_at_maximum, WithinAbs(-1.0, 3.0e-12));
            CHECK_THAT(result->limiting_point_determinant_at_step, WithinAbs(0.2, 3.0e-12));
            REQUIRE(result->limiting_rule >= 0);
            REQUIRE(static_cast<std::size_t>(result->limiting_rule) < rules.size());
            CHECK(result->limiting_element == rules[result->limiting_rule].element);
            CHECK(result->limiting_point >= 0);
            CHECK(result->limiting_point < rules[result->limiting_rule].integration_rule->GetNPoints());
        }
    }
}

TEST_CASE("Geometry step honors a positive determinant floor", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture;
    const mfem::Vector accepted = fixture.zero_displacement();
    const mfem::Vector direction = fixture.project_direction(compress_x);
    const auto rules = fixture.sample_rules();
    const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, accepted, direction, rules, {.determinant_floor = 0.25, .fraction_to_boundary = 0.8});

    REQUIRE(result);
    CHECK(result->limited_by_geometry);
    CHECK_THAT(result->boundary_step, WithinAbs(0.375, 2.0e-13));
    CHECK_THAT(result->step, WithinAbs(0.3, 2.0e-13));
    CHECK(result->limiting_point_determinant_at_step > 0.25);
}

TEST_CASE("Geometry step leaves an unconstrained step unchanged", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture;
    const mfem::Vector accepted = fixture.zero_displacement();
    const mfem::Vector direction = fixture.project_direction(expand_x);
    const auto rules = fixture.sample_rules();
    const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, accepted, direction, rules);

    REQUIRE(result);
    CHECK_FALSE(result->limited_by_geometry);
    CHECK_THAT(result->boundary_step, WithinAbs(1.0, 2.0e-13));
    CHECK_THAT(result->step, WithinAbs(1.0, 2.0e-13));
    CHECK_THAT(result->minimum_determinant_at_maximum, WithinAbs(2.0, 2.0e-13));
    CHECK_THAT(result->limiting_point_determinant_at_step, WithinAbs(2.0, 2.0e-13));
    CHECK(result->limiting_element == -1);
    CHECK(result->limiting_rule == -1);
    CHECK(result->limiting_point == -1);
}

TEST_CASE("Geometry step rescales the polynomial to the requested maximum step", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture;
    const mfem::Vector zero = fixture.zero_displacement();
    const mfem::Vector direction = fixture.project_direction(compress_x);
    const auto rules = fixture.sample_rules();
    const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, zero, direction, rules, {.maximum_step = 0.25});
    REQUIRE(result);
    CHECK_FALSE(result->limited_by_geometry);
    CHECK(result->step == 0.25);
    CHECK_THAT(result->minimum_determinant_at_maximum, WithinAbs(0.5, 2.0e-13));

    const auto endpoint = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, zero, direction, rules, {.maximum_step = 0.5});
    REQUIRE(endpoint);
    CHECK(endpoint->limited_by_geometry);
    CHECK_THAT(endpoint->step, WithinAbs(0.45, 2.0e-13));
}

TEST_CASE("Geometry step rejects invalid inputs as values", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture;
    mfem::Vector zero = fixture.zero_displacement();
    auto rules = fixture.sample_rules();
    const auto estimate = [&](const deformation::GeometryStepOptions options = {}) {
        return backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, zero, zero, rules, options);
    };

    SECTION("Invalid numerical controls") {
        CHECK_FALSE(estimate({.maximum_step = 0.0}));
        CHECK_FALSE(estimate({.maximum_step = std::numeric_limits<double>::infinity()}));
        CHECK_FALSE(estimate({.determinant_floor = -1.0}));
        CHECK_FALSE(estimate({.fraction_to_boundary = 0.0}));
        CHECK_FALSE(estimate({.fraction_to_boundary = 1.0}));
        CHECK_FALSE(estimate({.fraction_to_boundary = std::numeric_limits<double>::quiet_NaN()}));
    }
    SECTION("Empty or invalid rules") {
        rules.clear();
        CHECK_FALSE(estimate());
        rules.push_back({.element = 0, .integration_rule = nullptr});
        CHECK_FALSE(estimate());
        rules = fixture.sample_rules();
        rules[0].element = fixture.mesh.GetNE();
        CHECK_FALSE(estimate());
    }
    SECTION("Invalid coefficients") {
        zero.SetSize(1);
        CHECK_FALSE(estimate());
        zero = fixture.zero_displacement();
        zero(0) = std::numeric_limits<double>::quiet_NaN();
        CHECK_FALSE(estimate());
        zero = fixture.zero_displacement();
        fixture.compactification_coordinate(0) = std::numeric_limits<double>::infinity();
        CHECK_FALSE(estimate());
    }
    SECTION("Folded accepted state") {
        const mfem::Vector invalid_accepted = fixture.project_direction(compress_x);
        const auto result = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, invalid_accepted, zero, rules);
        REQUIRE_FALSE(result);
        CHECK(result.error().reason == deformation::DeformationFailureReason::invalid_geometry);
        CHECK(result.error().element >= 0);
        CHECK(result.error().point >= 0);
    }
    SECTION("Accepted state must lie strictly above the floor") {
        const auto result = estimate({.determinant_floor = 1.0});
        REQUIRE_FALSE(result);
        CHECK(result.error().reason == deformation::DeformationFailureReason::invalid_geometry);
    }
}

TEST_CASE("Geometry step rejects a non-affine exterior model before evaluation", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture(true);
    mapping::DomainMapper mapper({.dimension = dimension, .vacuum_element_attribute_ID = 2}, std::make_unique<NonAffineExteriorMap>());
    const mfem::Vector zero = fixture.zero_displacement();
    const auto rules = fixture.sample_rules();
    const auto result = backend::estimate_geometry_step_limit(mapper, fixture.displacement_space, fixture.compactification_coordinate, zero, zero, rules);
    REQUIRE_FALSE(result);
    CHECK(result.error().reason == deformation::DeformationFailureReason::invalid_configuration);
}

TEST_CASE("Geometry step retains the union of nonnested sampling rules", "[unit][deformation][geometry_step]") {
    GeometryFixture fixture(false, mfem::Ordering::byVDIM, 3);
    const mfem::Vector zero = fixture.zero_displacement();
    const mfem::Vector direction = fixture.project_direction([](const mfem::Vector& position, mfem::Vector& value) {
        value.SetSize(dimension);
        value = 0.0;
        const double centered_x = position(0) - 0.5;
        value(0) = -2.0 * position(0) + 8.0 * centered_x * centered_x * centered_x;
    });
    const std::array rules{
        backend::GeometrySampleRule{.element = 0, .integration_rule = &mfem::IntRules.Get(mfem::Geometry::CUBE, 1)},
        backend::GeometrySampleRule{.element = 0, .integration_rule = &mfem::IntRules.Get(mfem::Geometry::CUBE, 3)}
    };

    // The cubic direction compresses at the central point, while its x derivative vanishes at the two-point Gauss abscissae.
    const auto highest_only = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, zero, direction, std::span(rules).subspan(1));
    const auto both = backend::estimate_geometry_step_limit(fixture.mapper, fixture.displacement_space, fixture.compactification_coordinate, zero, direction, rules);
    REQUIRE(highest_only);
    REQUIRE(both);
    CHECK_FALSE(highest_only->limited_by_geometry);
    CHECK(both->limited_by_geometry);
    CHECK_THAT(both->boundary_step, WithinAbs(0.5, 2.0e-13));
    CHECK(both->limiting_rule == 0);
    CHECK(both->sampled_point_count == static_cast<std::size_t>(rules[0].integration_rule->GetNPoints() + rules[1].integration_rule->GetNPoints()));
}

TEST_CASE("Determinant polynomial preserves first roots and dimensional degree", "[unit][deformation][geometry_step]") {
    using deformation::detail::DeterminantPolynomial;
    using deformation::detail::JacobianEntries;
    using deformation::detail::determinant_polynomial;
    using deformation::detail::evaluate_polynomial;
    using deformation::detail::first_boundary_parameter;

    JacobianEntries base{};
    JacobianEntries direction{};
    for (int component = 0; component < 3; ++component) {
        base[component][component] = 1.0;
        direction[component][component] = -2.0;
    }
    for (int spatial_dimension = 1; spatial_dimension <= 3; ++spatial_dimension) {
        const DeterminantPolynomial coefficients = determinant_polynomial(base, direction, spatial_dimension, 0.75, 0.125);
        for (const double parameter : {0.0, 0.2, 0.6, 1.0}) {
            const double expected = std::pow(1.0 - 1.5 * parameter, spatial_dimension) - 0.125;
            CHECK_THAT(evaluate_polynomial(coefficients, parameter), WithinAbs(expected, 2.0e-15));
        }
    }

    CHECK(std::isinf(first_boundary_parameter({1.0, 0.0, 0.0, 0.0})));
    CHECK(std::isinf(first_boundary_parameter({1.0, 2.0, 3.0, 4.0})));
    CHECK_THAT(first_boundary_parameter({1.0, -2.0, 0.0, 0.0}), WithinAbs(0.5, 2.0e-13));
    CHECK_THAT(first_boundary_parameter({1.0, -4.0, 4.0, 0.0}), WithinAbs(0.5, 2.0e-13));
    CHECK_THAT(first_boundary_parameter({1.0, -5.5, 8.5, -3.0}), WithinAbs(1.0 / 3.0, 2.0e-13));
    CHECK_THAT(first_boundary_parameter({1.0e300, -4.0e300, 4.0e300, 0.0}), WithinAbs(0.5, 2.0e-13));
    CHECK_THAT(first_boundary_parameter({1.0e300, -5.5e300, 8.5e300, -3.0e300}), WithinAbs(1.0 / 3.0, 2.0e-13));
}