module;
#include <mfem.hpp>

export module mean_field:physics.gravity;

export import :fem;

export namespace mean_field::physics {
    struct GravitySolution {
        mfem::ParGridFunction gradPhi;
        mfem::ParGridFunction phi;

        explicit GravitySolution(fem::FEM &fem)
            : gradPhi(fem.gravityFluxFes.get()),
              phi(fem.gravityPotentialFes.get()) {
        }
    };

    GravitySolution grav_potential(
        fem::FEM &f,
        const utils::Args &args,
        const mfem::GridFunction &rho,
        bool phi_warm = false
    );

    GravitySolution grav_potential_new(
        fem::FEM &f,
        const utils::Args &args,
        const mfem::GridFunction &rho,
        const mfem::GridFunction &displacement
    );

    mfem::GridFunction get_potential(
        fem::FEM &fem,
        const utils::Args &args,
        const mfem::GridFunction &rho,
        bool warm = false
    );

    mfem::DenseMatrix compute_quadrupole_moment_tensor(
        const fem::FEM &fem,
        const mfem::GridFunction &rho,
        const mfem::Vector &com
    );

    double l2_multipole_potential(
        const fem::FEM &fem,
        double total_mass,
        const mfem::Vector &phys_x
    );

    void update_stiffness_matrix(fem::FEM &fem);
} // namespace mean_field::physics