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

export module mean_field:operators.gravity_field;
export import :fem;
export import :mapping.domain_mapper;
export import :operators.gravity_field_jacobian;
export import :operators.context.gravity_field;

export namespace mean_field::operators {
    enum class GravityResidualBlock : std::uint8_t { gradient_equation = 0, poisson_equation = 1, count = 2 };

    constexpr int gravity_residual_block_index(const GravityResidualBlock block) noexcept {
        return static_cast<int>(block);
    }

    inline constexpr int gravity_residual_block_count = gravity_residual_block_index(GravityResidualBlock::count);

    class GravityFieldOperator final : public mfem::Operator {
    public:
        GravityFieldOperator(
            fem::FEM &f,
            const mapping::DomainMapperStateless &domain_mapper,
            context::gravity_field::GravityFieldLinearizationContext &linearization_context,
            const mfem::Array<int> &state_true_offsets,
            GravityFieldJacobianOperator &jacobian
        );

        context::gravity_field::GravityFieldPreparationReport Prepare(
            const mfem::Vector &state,
            const context::gravity_field::GravityFieldRevisions &revisions
        );

        void Mult(
            const mfem::Vector &state,
            mfem::Vector &residual
        ) const override;

        Operator &GetGradient(const mfem::Vector &state) const override;

        [[nodiscard]] const mfem::Array<int> &GetStateTrueOffsets() const noexcept;

        [[nodiscard]] const mfem::Array<int> &GetResidualTrueOffsets() const noexcept;

        [[nodiscard]] context::gravity_field::GravityFieldLinearizationContext &GetLinearizationContext() noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetLinearizationContext() const noexcept;

        void ApplyGravityUnknowns(
            const mfem::Vector &gravity_gradient,
            const mfem::Vector &gravity_potential,
            const context::gravity_field::GravityFieldGeometryContext &geometry_context,
            mfem::Vector &action
        ) const;

        void ApplyDensitySource(
            const mfem::Vector &density,
            const context::gravity_field::GravityFieldGeometryContext &geometry_context,
            mfem::Vector &action
        ) const;

    private:
        fem::FEM &m_fem;
        const mapping::DomainMapperStateless &m_domain_mapper;
        context::gravity_field::GravityFieldLinearizationContext &m_linearization_context;
        mfem::Array<int> m_state_true_offsets;
        mfem::Array<int> m_residual_true_offsets;
        GravityFieldJacobianOperator &m_jacobian;
    };

    class ReducedGravityFieldOperator final : public mfem::Operator {
    public:
        ReducedGravityFieldOperator(
            GravityFieldOperator &gravity_field_operator,
            context::gravity_field::GravityFieldGeometryContext &gravity_field_geometry_context,
            const mfem::Vector &displacement
        );

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

        void SetDisplacement(const mfem::Vector &displacement);

        [[nodiscard]] const mfem::Vector &GetDisplacement() const;

        void BuildRightHandSide(
            const mfem::Vector &density,
            mfem::Vector &right_hand_side
        ) const;

        void Mult(
            const mfem::Vector &gravity_state,
            mfem::Vector &action
        ) const override;

        [[nodiscard]] GravityFieldOperator &GetGravityFieldOperator() noexcept;

        [[nodiscard]] const GravityFieldOperator &GetGravityFieldOperator() const noexcept;

        [[nodiscard]] context::gravity_field::GravityFieldGeometryContext &GetGeometryContext() noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldGeometryContext &GetGeometryContext() const noexcept;

        [[nodiscard]] const mfem::Array<int> &GetGravityTrueOffsets() const noexcept;

    private:
        void ValidateDisplacement(const mfem::Vector &displacement) const;

        void ValidateDensity(const mfem::Vector &density) const;

        void ValidateGravityState(const mfem::Vector &gravity_state) const;

    private:
        GravityFieldOperator &m_gravity_field_operator;
        mfem::Array<int> m_gravity_true_offsets;
        context::gravity_field::GravityFieldGeometryContext &m_gravity_field_geometry_context;
    };
} // namespace mean_field::operators