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

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 refreshed_variation_state{false};

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

    class GravityFieldGeometryContext {
    public:
        GravityFieldGeometryContext(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &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]] const PreparedMappedHDivMassOperator &GetMassOperator() const;
        [[nodiscard]] const PreparedMappedGravitySourceOperator &GetSourceOperator() 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:
        const fem::FEM &m_fem;
        const mapping::DomainMapperStateless &m_domain_mapper;

        std::unique_ptr<PreparedMappedHDivMassOperator> m_mass_operator;
        std::unique_ptr<PreparedMappedGravitySourceOperator> m_source_operator;

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

        DiscretizationRevision m_discretization_revision;
        DisplacementRevision m_displacement_revision;

        bool m_is_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::DomainMapperStateless &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]] 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