#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include <cstdint>
#include <memory>
#include <mfem.hpp>
#include <mpi.h>
#include <stdexcept>
#include <utility>
#include <vector>

import mean_field;

namespace {
    using Catch::Matchers::WithinAbs;

    constexpr int dimension = 3;

    [[nodiscard]] mfem::Mesh make_serial_mesh(const int attribute) {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON, 2.0, 1.0, 1.0);
        for (int element = 0; element < mesh.GetNE(); ++element) {
            mesh.GetElement(element)->SetAttribute(attribute);
        }
        return mesh;
    }

    [[nodiscard]] std::unique_ptr<const mean_field::mapping::compactification::ExteriorDomainMap>
    make_kelvin_compactification() {
        return std::make_unique<mean_field::mapping::compactification::KelvinCompactification>(
            mean_field::mapping::compactification::options::KelvinCompactificationOptions{
                .r_star_ref = 1.0, .r_inf_ref = 4.0
            }
        );
    }

    struct GeometryFixture final {
        mfem::Mesh serialMesh;
        mfem::ParMesh mesh;
        mfem::H1_FECollection displacementCollection;
        mfem::ParFiniteElementSpace displacementSpace;
        mfem::H1_FECollection compactificationCollection;
        mfem::ParFiniteElementSpace compactificationSpace;
        mfem::ParGridFunction compactificationCoordinate;
        mean_field::mapping::DomainMapper mapper;

        explicit GeometryFixture(const bool compactified = false)
            : serialMesh(make_serial_mesh(compactified ? 2 : 1)),
              mesh(
                  MPI_COMM_WORLD,
                  serialMesh
              ),
              displacementCollection(
                  1,
                  dimension
              ),
              displacementSpace(
                  &mesh,
                  &displacementCollection,
                  dimension,
                  mfem::Ordering::byNODES
              ),
              compactificationCollection(
                  1,
                  dimension
              ),
              compactificationSpace(
                  &mesh,
                  &compactificationCollection
              ),
              compactificationCoordinate(&compactificationSpace),
              mapper(
                  {.dimension                = dimension,
                   .vacuum_element_attribute = 2},
                  make_kelvin_compactification()
              ) {
            compactificationCoordinate = 0.0;
        }

        [[nodiscard]] mfem::Vector zero_true_vector() const {
            mfem::Vector result(displacementSpace.GetTrueVSize());
            result = 0.0;
            return result;
        }

        template <typename Function> [[nodiscard]] mfem::Vector project_direction(Function &&function) {
            mfem::VectorFunctionCoefficient coefficient(dimension, std::forward<Function>(function));
            mfem::ParGridFunction field(&displacementSpace);
            field.ProjectCoefficient(coefficient);
            mfem::Vector result;
            field.GetTrueDofs(result);
            return result;
        }

        [[nodiscard]] std::vector<mean_field::deformation::NewtonStepGeometryRule> geometry_rules() {
            std::vector<mean_field::deformation::NewtonStepGeometryRule> result;
            result.reserve(static_cast<std::size_t>(mesh.GetNE()));
            for (int element = 0; element < mesh.GetNE(); ++element) {
                mfem::ElementTransformation *transformation = mesh.GetElementTransformation(element);
                result.push_back(
                    {.element = element, .integrationRule = &mfem::IntRules.Get(transformation->GetGeometryType(), 2)}
                );
            }
            return result;
        }
    };

    void compress_x(
        const mfem::Vector &position,
        mfem::Vector &value
    ) {
        value.SetSize(dimension);
        value    = 0.0;
        value(0) = -2.0 * position(0);
    }

    void compress_x_and_y(
        const mfem::Vector &position,
        mfem::Vector &value
    ) {
        value.SetSize(dimension);
        value    = 0.0;
        value(0) = -2.0 * position(0);
        value(1) = -2.0 * position(1);
    }

    void expand_x(
        const mfem::Vector &position,
        mfem::Vector &value
    ) {
        value.SetSize(dimension);
        value    = 0.0;
        value(0) = position(0);
    }
} // namespace

TEST_CASE(
    "Safe Newton Step Finds The First Mapping Boundary",
    "[deformation][newton][geometry][mpi]"
) {
    GeometryFixture fixture;
    const mfem::Vector accepted  = fixture.zero_true_vector();
    const mfem::Vector direction = fixture.project_direction(compress_x);
    const auto rules             = fixture.geometry_rules();

    const auto estimate          = mean_field::deformation::estimate_largest_safe_newton_step_size(
        fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, accepted, direction, rules,
        {.maximumStepSize = 1.0, .determinantFloor = 0.0, .fractionToBoundarySafety = 0.8}
    );

    CHECK(estimate.limitedByGeometry);
    CHECK_THAT(estimate.boundaryStepSize, WithinAbs(0.5, 2.0e-13));
    CHECK_THAT(estimate.stepSize, WithinAbs(0.4, 2.0e-13));
    CHECK_THAT(estimate.minimumDeterminantAtAcceptedState, WithinAbs(1.0, 2.0e-13));
    CHECK_THAT(estimate.minimumDeterminantAtMaximumStepSize, WithinAbs(-1.0, 2.0e-13));
    CHECK_THAT(estimate.limitingPointDeterminantAtStepSize, WithinAbs(0.2, 2.0e-13));
    CHECK(estimate.sampledQuadraturePointCount > 0);
    CHECK(estimate.limitingRank == 0);
    CHECK(estimate.limitingElement >= 0);
    CHECK(estimate.limitingRule >= 0);
    CHECK(estimate.limitingQuadraturePoint >= 0);
}

TEST_CASE(
    "Safe Newton Step Detects A Tangent Singularity Before An Admissible Endpoint",
    "[deformation][newton][geometry][mpi]"
) {
    GeometryFixture fixture(true);
    const mfem::Vector accepted  = fixture.zero_true_vector();
    const mfem::Vector direction = fixture.project_direction(compress_x_and_y);
    const auto rules             = fixture.geometry_rules();

    const auto estimate          = mean_field::deformation::estimate_largest_safe_newton_step_size(
        fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, accepted, direction, rules
    );

    // det(J(alpha)) = (1 - 2 alpha)^2.  Both endpoints are positive;
    // checking only alpha=1 would miss the singularity at alpha=1/2.
    CHECK(estimate.limitedByGeometry);
    CHECK_THAT(estimate.minimumDeterminantAtMaximumStepSize, WithinAbs(1.0, 3.0e-13));
    CHECK_THAT(estimate.boundaryStepSize, WithinAbs(0.5, 3.0e-13));
    CHECK_THAT(estimate.stepSize, WithinAbs(0.45, 3.0e-13));
    CHECK_THAT(estimate.limitingPointDeterminantAtStepSize, WithinAbs(0.01, 3.0e-13));
}

TEST_CASE(
    "Safe Newton Step Honors A Positive Determinant Floor",
    "[deformation][newton][geometry][mpi]"
) {
    GeometryFixture fixture;
    const mfem::Vector accepted  = fixture.zero_true_vector();
    const mfem::Vector direction = fixture.project_direction(compress_x);
    const auto rules             = fixture.geometry_rules();

    const auto estimate          = mean_field::deformation::estimate_largest_safe_newton_step_size(
        fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, accepted, direction, rules,
        {.maximumStepSize = 1.0, .determinantFloor = 0.25, .fractionToBoundarySafety = 0.8}
    );

    CHECK(estimate.limitedByGeometry);
    CHECK_THAT(estimate.boundaryStepSize, WithinAbs(0.375, 2.0e-13));
    CHECK_THAT(estimate.stepSize, WithinAbs(0.3, 2.0e-13));
    CHECK(estimate.limitingPointDeterminantAtStepSize > 0.25);
}

TEST_CASE(
    "Safe Newton Step Leaves An Unconstrained Step Unchanged",
    "[deformation][newton][geometry][mpi]"
) {
    GeometryFixture fixture;
    const mfem::Vector accepted  = fixture.zero_true_vector();
    const mfem::Vector direction = fixture.project_direction(expand_x);
    const auto rules             = fixture.geometry_rules();

    const auto estimate          = mean_field::deformation::estimate_largest_safe_newton_step_size(
        fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, accepted, direction, rules
    );

    CHECK_FALSE(estimate.limitedByGeometry);
    CHECK_THAT(estimate.boundaryStepSize, WithinAbs(1.0, 2.0e-13));
    CHECK_THAT(estimate.stepSize, WithinAbs(1.0, 2.0e-13));
    CHECK_THAT(estimate.minimumDeterminantAtMaximumStepSize, WithinAbs(2.0, 2.0e-13));
    CHECK_THAT(estimate.limitingPointDeterminantAtStepSize, WithinAbs(2.0, 2.0e-13));
    CHECK(estimate.limitingRank == -1);
    CHECK(estimate.limitingElement == -1);
    CHECK(estimate.limitingRule == -1);
    CHECK(estimate.limitingQuadraturePoint == -1);
}

TEST_CASE(
    "Safe Newton Step Rejects Invalid Inputs Collectively",
    "[deformation][newton][geometry][mpi]"
) {
    GeometryFixture fixture;
    const mfem::Vector zero = fixture.zero_true_vector();
    const auto rules        = fixture.geometry_rules();

    CHECK_THROWS_AS(
        mean_field::deformation::estimate_largest_safe_newton_step_size(
            fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, zero, zero, rules,
            {.maximumStepSize = 0.0}
        ),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        mean_field::deformation::estimate_largest_safe_newton_step_size(
            fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, zero, zero, {}
        ),
        std::invalid_argument
    );

    const mfem::Vector invalidAccepted = fixture.project_direction(compress_x);
    CHECK_THROWS_AS(
        mean_field::deformation::estimate_largest_safe_newton_step_size(
            fixture.mapper, fixture.displacementSpace, fixture.compactificationCoordinate, invalidAccepted, zero, rules
        ),
        std::domain_error
    );
}