module;

#include <cstdint>
#include <limits>
#include <span>

#include <mfem.hpp>

export module mean_field:deformation.safe_newton_step;

export import :mapping.domain_mapper;

export namespace mean_field::deformation {
    /*
     * One element-local quadrature rule at which a downstream operator will
     * evaluate the mapped geometry.  A caller may provide several entries for
     * one element when several operators use different, non-nested rules.
     * The integration-rule object must outlive the call.
     */
    struct NewtonStepGeometryRule final {
        int element{-1};
        const mfem::IntegrationRule *integrationRule{nullptr};
    };

    struct LargestSafeNewtonStepSizeOptions final {
        double maximumStepSize{1.0};
        double determinantFloor{0.0};
        double fractionToBoundarySafety{0.9};

        void Validate() const;
    };

    /*
     * The estimated boundary is the first alpha in [0, maximumStepSize] at
     * which any sampled mapping determinant reaches determinantFloor.  When
     * no such point exists, boundaryStepSize equals maximumStepSize and
     * limitedByGeometry is false.  stepSize is the boundary multiplied by the
     * safety fraction only when geometry is limiting.
     *
     * Element and rule indices are local to limitingRank.  limitingRule is an
     * index into that rank's input span.
     */
    struct LargestSafeNewtonStepSizeEstimate final {
        double stepSize{0.0};
        double boundaryStepSize{0.0};
        double minimumDeterminantAtAcceptedState{std::numeric_limits<double>::quiet_NaN()};
        double minimumDeterminantAtMaximumStepSize{std::numeric_limits<double>::quiet_NaN()};
        double limitingPointDeterminantAtStepSize{std::numeric_limits<double>::quiet_NaN()};
        std::uint64_t sampledQuadraturePointCount{0};
        bool limitedByGeometry{false};
        int limitingRank{-1};
        int limitingElement{-1};
        int limitingRule{-1};
        int limitingQuadraturePoint{-1};
    };

    /*
     * Estimate the largest safe alpha for
     *
     *     displacement(alpha) = acceptedVolumeDisplacement
     *                           + alpha * volumeNewtonDirection.
     *
     * Both vectors use the displacement space's true-DOF layout.  The
     * compactification coordinate is held fixed.  The result is collective on
     * displacementSpace.GetComm() and is identical on every rank.
     *
     * The calculation is exact for the current domain mapper: its mapping
     * Jacobian is affine along a displacement direction, so each sampled
     * determinant is a polynomial of degree at most the spatial dimension.
     * Compactified elements require an exterior map that explicitly advertises
     * the same affine contract.
     */
    [[nodiscard]] LargestSafeNewtonStepSizeEstimate estimate_largest_safe_newton_step_size(
        const mapping::DomainMapper &domainMapper,
        const mfem::ParFiniteElementSpace &displacementSpace,
        const mfem::ParGridFunction &compactificationCoordinate,
        const mfem::Vector &acceptedVolumeDisplacement,
        const mfem::Vector &volumeNewtonDirection,
        std::span<const NewtonStepGeometryRule> geometryRules,
        const LargestSafeNewtonStepSizeOptions &options = {}
    );
} // namespace mean_field::deformation