MeanField/libmeanfield/interface/operators/contexts/gravity_field_context.cppm

module;
#include <compare>
#include <cstdint>
#include <expected>
#include <memory>
#include <mfem.hpp>
#include <stdexcept>

export module mean_field:operators.context.gravity_field;
export import :fem;
export import :field.mfem;
export import :mapping.domain_mapper;
export import :operators.prepared_gravity_source;
export import :operators.prepared_hdiv_mass;

export namespace mean_field::operators::context::gravity_field {
    template <typename Tag> struct Revision {
        std::uint64_t value{0};

        constexpr auto operator<=>(const Revision &) const = default;
    };

    struct DiscretizationRevisionTag { };
    struct DisplacementRevisionTag { };
    struct DensityRevisionTag { };
    struct GravityGradientRevisionTag { };
    struct GravityPotentialRevisionTag { };

    using DiscretizationRevision   = Revision<DiscretizationRevisionTag>;
    using DisplacementRevision     = Revision<DisplacementRevisionTag>;
    using DensityRevision          = Revision<DensityRevisionTag>;
    using GravityGradientRevision  = Revision<GravityGradientRevisionTag>;
    using GravityPotentialRevision = Revision<GravityPotentialRevisionTag>;

    struct GravityFieldRevisions {
        DiscretizationRevision discretization;
        DisplacementRevision displacement;
        DensityRevision density;
        GravityGradientRevision gravity_gradient;
        GravityPotentialRevision gravity_potential;
    };

    struct GravityFieldStateView {
        const mfem::Vector &density;
        const mfem::Vector &displacement;
        const mfem::Vector &gravity_gradient;
        const mfem::Vector &gravity_potential;
    };

    struct GravityFieldGeometryPreparation {
        bool reconstructed_operators{false};
        bool rebuilt_mass_operator{false};
        bool rebuilt_source_operator{false};
        bool rebuilt_divergence_operator{false};
        bool refreshed_variation_state{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return reconstructed_operators || rebuilt_mass_operator || rebuilt_source_operator ||
                   rebuilt_divergence_operator || refreshed_variation_state;
        }
    };

    enum class GravityFieldPreparationRejectionReason : std::uint8_t { invalid_mapping, non_finite_arithmetic };

    struct GravityFieldPreparationRejection final {
        GravityFieldPreparationRejectionReason reason{GravityFieldPreparationRejectionReason::invalid_mapping};
        mapping::MappingStatus mappingStatus{mapping::MappingStatus::valid};
    };

    template <typename Report>
    using GravityFieldPreparationResult = std::expected<Report, GravityFieldPreparationRejection>;

    [[noreturn]] inline void throwGravityFieldPreparationRejection(const GravityFieldPreparationRejection &rejection) {
        if (rejection.reason == GravityFieldPreparationRejectionReason::non_finite_arithmetic) {
            throw std::domain_error("Prepared gravity-field data contained non-finite arithmetic.");
        }
        throw std::domain_error("Prepared gravity-field data could not map the candidate geometry.");
    }

    class GravityFieldGeometryContext {
    public:
        GravityFieldGeometryContext(
            const fem::FEM &f,
            const mapping::DomainMapper &domain_mapper
        );

        GravityFieldGeometryContext(const GravityFieldGeometryContext &)            = delete;
        GravityFieldGeometryContext &operator=(const GravityFieldGeometryContext &) = delete;
        GravityFieldGeometryContext(GravityFieldGeometryContext &&)                 = delete;
        GravityFieldGeometryContext &operator=(GravityFieldGeometryContext &&)      = delete;

        GravityFieldGeometryPreparation Prepare(
            const mfem::Vector &displacement,
            DiscretizationRevision discretization_revision,
            DisplacementRevision displacement_revision
        );

        [[nodiscard]] GravityFieldPreparationResult<GravityFieldGeometryPreparation> TryPrepare(
            const mfem::Vector &displacement,
            DiscretizationRevision discretization_revision,
            DisplacementRevision displacement_revision
        );

        GravityFieldGeometryPreparation PreparePrimal(
            const mfem::Vector &displacement,
            DiscretizationRevision discretization_revision,
            DisplacementRevision displacement_revision
        );

        [[nodiscard]] GravityFieldPreparationResult<GravityFieldGeometryPreparation> TryPreparePrimal(
            const mfem::Vector &displacement,
            DiscretizationRevision discretization_revision,
            DisplacementRevision displacement_revision
        );

        [[nodiscard]] const PreparedMappedHDivMassOperator &GetMassOperator() const;
        [[nodiscard]] const PreparedMappedGravitySourceOperator &GetSourceOperator() const;
        [[nodiscard]] const mfem::Operator &GetDivergenceOperator() const;
        [[nodiscard]] const mfem::Operator &GetTransposeDivergenceOperator() const;
        [[nodiscard]] const mfem::Vector &GetDisplacementTrue() const;
        [[nodiscard]] const field::FieldDofMap &GetDisplacementMap() const noexcept;
        [[nodiscard]] DiscretizationRevision GetDiscretizationRevision() const noexcept;
        [[nodiscard]] DisplacementRevision GetDisplacementRevision() const noexcept;
        [[nodiscard]] bool IsPrepared() const noexcept;

    private:
        enum class PreparationMode : std::uint8_t { primal, linearization };

        [[nodiscard]] GravityFieldPreparationResult<GravityFieldGeometryPreparation> TryPrepareImpl(
            const mfem::Vector &displacement,
            DiscretizationRevision discretization_revision,
            DisplacementRevision displacement_revision,
            PreparationMode mode
        );

        const fem::FEM &m_fem;
        const mapping::DomainMapper &m_domain_mapper;

        std::unique_ptr<PreparedMappedHDivMassOperator> m_mass_operator;
        std::unique_ptr<PreparedMappedGravitySourceOperator> m_source_operator;
        std::unique_ptr<mfem::Operator> m_divergence_operator;
        std::unique_ptr<mfem::TransposeOperator> m_transpose_divergence_operator;

        field::FieldDofMap m_displacement_map;
        mfem::Vector m_displacement_true;

        DiscretizationRevision m_discretization_revision;
        DisplacementRevision m_displacement_revision;

        bool m_is_prepared{false};
        bool m_variation_state_prepared{false};
    };

    struct GravityFieldPreparationReport {
        GravityFieldGeometryPreparation geometry;
        bool updated_density{false};
        bool updated_gravity_gradient{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return geometry.DidAnyWork() || updated_density || updated_gravity_gradient;
        }
    };

    class GravityFieldLinearizationContext {
    public:
        GravityFieldLinearizationContext(
            const fem::FEM &f,
            const mapping::DomainMapper &domain_mapper
        );

        GravityFieldLinearizationContext(const GravityFieldLinearizationContext &)            = delete;
        GravityFieldLinearizationContext &operator=(const GravityFieldLinearizationContext &) = delete;
        GravityFieldLinearizationContext(GravityFieldLinearizationContext &&)                 = delete;
        GravityFieldLinearizationContext &operator=(GravityFieldLinearizationContext &&)      = delete;

        GravityFieldPreparationReport Prepare(
            const GravityFieldStateView &state,
            const GravityFieldRevisions &revisions
        );

        [[nodiscard]] GravityFieldPreparationResult<GravityFieldPreparationReport> TryPrepare(
            const GravityFieldStateView &state,
            const GravityFieldRevisions &revisions
        );

        [[nodiscard]] const GravityFieldGeometryContext &GetGeometryContext() const;
        [[nodiscard]] const mfem::Vector &GetDensityTrue() const;
        [[nodiscard]] const mfem::Vector &GetGravityGradientTrue() const;
        [[nodiscard]] const field::FieldDofMap &GetDensityMap() const noexcept;
        [[nodiscard]] const field::FieldDofMap &GetDisplacementMap() const noexcept;
        [[nodiscard]] const field::FieldDofMap &GetGravityGradientMap() const noexcept;
        [[nodiscard]] const field::FieldDofMap &GetGravityPotentialMap() const noexcept;
        [[nodiscard]] const GravityFieldRevisions &GetRevisions() const;
        [[nodiscard]] bool IsPrepared() const noexcept;

    private:
        const fem::FEM &m_fem;

        GravityFieldGeometryContext m_geometry_context;

        field::FieldDofMap m_density_map;
        field::FieldDofMap m_gravity_gradient_map;
        field::FieldDofMap m_gravity_potential_map;
        mfem::Vector m_density_true;
        mfem::Vector m_gravity_gradient_true;

        GravityFieldRevisions m_revisions;
        bool m_is_prepared{false};
    };
} // namespace mean_field::operators::context::gravity_field