FEM setup_fem(const std::string& filename, const bool verbose) {
    FEM fem;
    fem.mesh = std::make_unique<mfem::Mesh>(filename, 0, 0);
    fem.mesh->EnsureNodes();
    int dim = fem.mesh->Dimension();
    fem.H1_fec   = std::make_unique<mfem::H1_FECollection>(2, dim);
    fem.H1_fes   = std::make_unique<mfem::FiniteElementSpace>(fem.mesh.get(), fem.H1_fec.get());
    fem.Vec_H1_fes = std::make_unique<mfem::FiniteElementSpace>(fem.mesh.get(), fem.H1_fec.get(), dim, mfem::Ordering::byNODES);

    fem.block_offsets.SetSize(3);
    fem.block_offsets[0] = 0;
    fem.block_offsets[1] = fem.H1_fes->GetTrueVSize();
    fem.block_offsets[2] = fem.H1_fes->GetTrueVSize() + fem.Vec_H1_fes->GetTrueVSize();
    fem.com.SetSize(dim); fem.com = 0.0;
    fem.Q.SetSize(dim, dim); fem.Q = 0.0;
    return fem;
}

void view_mesh(const std::string& host, int port, const mfem::Mesh& mesh, const mfem::GridFunction& gf, const std::string& title) {
    mfem::socketstream sol_sock(host.c_str(), port);
    if (!sol_sock.is_open()) return;
    sol_sock << "solution\n" << mesh << gf;
    sol_sock << "window_title '" << title << "'\n" << std::flush;

}

double domain_integrate_grid_function(const FEM& fem, const mfem::GridFunction& gf) {
    mfem::LinearForm lf(fem.H1_fes.get());
    mfem::GridFunctionCoefficient gf_c(&gf);
    lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(gf_c));
    lf.Assemble();
    return lf.Sum();
}

mfem::Vector get_com(const FEM& fem, const mfem::GridFunction &rho) {
    const int dim = fem.mesh->Dimension();
    mfem::Vector com(dim);
    com = 0.0;
    double total_mass = 0.0;

    for (int i = 0; i < fem.H1_fes->GetNE(); ++i) {
        mfem::ElementTransformation *trans = fem.H1_fes->GetElementTransformation(i);
        const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), fem.H1_fes->GetOrder(0) + trans->OrderW());

        for (int j = 0; j < ir.GetNPoints(); ++j) {
            const mfem::IntegrationPoint &ip = ir.IntPoint(j);
            trans->SetIntPoint(&ip);

            double weight = trans->Weight() * ip.weight;
            double rho_val = rho.GetValue(i, ip);

            mfem::Vector phys_point(dim);
            trans->Transform(ip, phys_point);

            const double mass_term = rho_val * weight;
            total_mass += mass_term;

            for (int d = 0; d < dim; ++d) {
                com(d) += phys_point(d) * mass_term;
            }
        }
    }
    com /= total_mass;
    return com;
}

double centrifugal_potential(const mfem::Vector& x, double omega) {
    const double s2 = std::pow(x(0), 2) + std::pow(x(1), 2);
    return -0.5 * s2 * std::pow(omega, 2);
}

double get_moment_of_inertia(const FEM& fem, const mfem::GridFunction& rho) {
    auto s2_func = [](const mfem::Vector& x) {
        return std::pow(x(0), 2) + std::pow(x(1), 2);
    };

    mfem::FunctionCoefficient s2_coeff(s2_func);
    mfem::GridFunctionCoefficient rho_coeff(&rho);

    mfem::ProductCoefficient I_integrand ( rho_coeff, s2_coeff );
    mfem::LinearForm I_lf(fem.H1_fes.get());
    I_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(I_integrand));
    I_lf.Assemble();

    return I_lf.Sum();
}

std::unique_ptr<mfem::GridFunction> grav_potential(const FEM& fem, const Args &args, const mfem::GridFunction& rho) {
    auto phi = std::make_unique<mfem::GridFunction>(fem.H1_fes.get());

    mfem::Array<int> ess_bdr(fem.mesh->bdr_attributes.Max());
    ess_bdr = 1;

    mfem::GridFunctionCoefficient rho_coeff(&rho);
    double total_mass = domain_integrate_grid_function(fem, rho);

    auto grav_potential = [&fem, &total_mass](const mfem::Vector& x) {
        return l2_multipole_potential(fem, total_mass, x);
    };

    mfem::FunctionCoefficient phi_bdr_coeff(grav_potential);
    mfem::Array<int> ess_tdof_list;
    fem.H1_fes->GetEssentialTrueDofs(ess_bdr, ess_tdof_list);

    auto laplacian = std::make_unique<mfem::BilinearForm>(fem.H1_fes.get());
    laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator());
    laplacian->Assemble();
    laplacian->Finalize();

    mfem::ConstantCoefficient four_pi_G(-4.0 * M_PI * G);
    mfem::ProductCoefficient rhs_coeff(rho_coeff, four_pi_G);
    mfem::LinearForm b(fem.H1_fes.get());
    b.AddDomainIntegrator(new mfem::DomainLFIntegrator(rhs_coeff));
    b.Assemble();

    mfem::OperatorPtr A;
    mfem::Vector B, X;
    laplacian->FormLinearSystem(ess_tdof_list, *phi, b, A, X, B);

    mfem::GSSmoother prec;
    mfem::CGSolver cg;
    cg.SetPreconditioner(prec);
    cg.SetOperator(*A);
    cg.SetRelTol(args.p.tol);
    cg.SetMaxIter(args.p.max_iters);
    cg.SetPrintLevel(0);

    cg.Mult(B, X);
    laplacian->RecoverFEMSolution(X, b, *phi);
    return phi;
}

std::unique_ptr<mfem::GridFunction> get_potential(const FEM &fem, const Args &args, const mfem::GridFunction &rho) {
    auto phi = grav_potential(fem, args, rho);
    mfem::GridFunctionCoefficient rho_coeff(&rho);

    if (args.r.enabled) {
        auto rot = [&args](const mfem::Vector& x) {
            return centrifugal_potential(x, args.r.omega);
        };

        mfem::FunctionCoefficient centrifugal_coeff(rot);
        mfem::GridFunction centrifugal_gf(fem.H1_fes.get());
        centrifugal_gf.ProjectCoefficient(centrifugal_coeff);
        (*phi) += centrifugal_gf;
    }

    return phi;
}

mfem::DenseMatrix compute_quadrupole_moment_tensor(const FEM& fem, const mfem::GridFunction& rho, const mfem::Vector& com) {
    const int dim = fem.mesh->Dimension();
    mfem::DenseMatrix Q(dim, dim);
    Q = 0.0;

    for (int i = 0; i < fem.H1_fes->GetNE(); ++i) {
        mfem::ElementTransformation *trans = fem.mesh->GetElementTransformation(i);
        const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), 2 * fem.H1_fes->GetOrder(0) + trans->OrderW());
        for (int j = 0; j < ir.GetNPoints(); ++j) {
            const mfem::IntegrationPoint &ip = ir.IntPoint(j);
            trans->SetIntPoint(&ip);

            const double weight = trans->Weight() * ip.weight;
            const double rho_val = rho.GetValue(i, ip);

            mfem::Vector phys_point(dim);
            trans->Transform(ip, phys_point);

            mfem::Vector x_prime(dim);
            double r_sq = 0.0;

            for (int d = 0; d < dim; ++d) {
                x_prime(d) = phys_point(d) - com(d);
                r_sq += x_prime(d) * x_prime(d);
            }

            for (int m = 0; m < dim; ++m) {
                for (int n = 0; n < dim; ++n) {
                    const double delta = (m == n) ? 1.0 : 0.0;
                    const double contrib = 3.0 * x_prime(m) * x_prime(n) - delta * r_sq;
                    Q(m, n) += rho_val * contrib * weight;
                }
            }
        }
    }
    return Q;
}

double l2_multipole_potential(const FEM &fem, const double total_mass, const mfem::Vector &x) {
    const double r = x.Norml2();
    if (r < 1e-12) return 0.0;

    const int dim = fem.mesh->Dimension();

    mfem::Vector n(x);
    n /= r;

    double l2_mult_factor = 0.0;
    for (int i = 0; i < dim; ++i) {
        for (int j = 0; j < dim; ++j) {
            l2_mult_factor += fem.Q(i, j) * n(i) * n(j);
        }
    }

    const double l2_contrib = (G / (2.0 * std::pow(r, 3))) * l2_mult_factor;

    const double l0_contrib = -G * total_mass / r;

    // l1 contribution is zero for a system centered on its COM
    return l0_contrib + l2_contrib;
}

void ConserveMass(const FEM& fem, mfem::GridFunction& rho, double target_mass) {
    const double current_mass = domain_integrate_grid_function(fem, rho);
    if (current_mass > 1e-15) rho *= (target_mass / current_mass);
}