module;
#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 {
    class GravityFieldOperator final : public mfem::Operator {
    public:
        GravityFieldOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domain_mapper,
            context::gravity_field::GravityFieldLinearizationContext &linearization_context,
            const mfem::Array<int> &state_offsets,
            GravityFieldJacobianOperator &jacobian
        );

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

        [[nodiscard]] context::gravity_field::GravityFieldPreparationResult<
            context::gravity_field::GravityFieldPreparationReport>
        TryPrepare(
            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> &GetStateOffsets() const noexcept;

        [[nodiscard]] const mfem::Array<int> &GetResidualOffsets() 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::DomainMapper &m_domain_mapper;
        context::gravity_field::GravityFieldLinearizationContext &m_linearization_context;
        mfem::Array<int> m_state_offsets;
        mfem::Array<int> m_residual_offsets;
        GravityFieldJacobianOperator &m_jacobian;

        mutable mfem::Vector m_potential_true;
        mutable mfem::Vector m_transpose_divergence_action_true;
        mutable mfem::Vector m_transpose_divergence_action;
        mutable mfem::Vector m_gradient_true;
        mutable mfem::Vector m_divergence_action_true;
    };

    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> &GetGravityOffsets() 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_offsets;
        context::gravity_field::GravityFieldGeometryContext &m_gravity_field_geometry_context;
        mfem::Vector m_displacement;
    };

    class ReducedGravityFieldPreconditioner final : public mfem::Solver {
    public:
        ReducedGravityFieldPreconditioner(
            const fem::FEM &f,
            const context::gravity_field::GravityFieldGeometryContext &geometry_context
        );

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

        void SetOperator(const mfem::Operator &gravity_operator) override;

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

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

    private:
        field::FieldDofMap m_flux_map;
        field::FieldDofMap m_potential_map;
        mfem::Array<int> m_offsets;
        mfem::Array<int> m_empty_tdofs;

        std::unique_ptr<mfem::OperatorJacobiSmoother> m_mass_preconditioner;
        std::unique_ptr<mfem::HypreParMatrix> m_schur;
        std::unique_ptr<mfem::HypreBoomerAMG> m_potential_preconditioner;

        mutable mfem::Vector m_potential_rhs_true;
        mutable mfem::Vector m_potential_action_true;
    };
} // namespace mean_field::operators