module;
#include "mean_field.h"

module mean_field;

namespace mean_field::physics {
    double compute_moment_of_inertia(const fem::FEM &fem, const mfem::GridFunction &rho_ref) {
        double local_I = 0.0;

        for (int i = 0; i < fem.mesh->GetNE(); i++) {
            if (fem.mesh->GetAttribute(i) == 3) continue;

            mfem::ElementTransformation *T = fem.mesh->GetElementTransformation(i);
            const mfem::IntegrationRule &ir = *fem.int_rule;

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

                const double rho_hat = rho_ref.GetValue(i, ip);

                mfem::Vector x_phys;
                fem.mapping->GetPhysicalPoint(*T, ip, x_phys);

                const double r_cyl_sq = x_phys(0) * x_phys(0) + x_phys(1) * x_phys(1);
                const double detJ = std::fabs(fem.mapping->ComputeDetJ(*T, ip));
                const double weight = T->Weight() * ip.weight * detJ;

                local_I += rho_hat * r_cyl_sq * weight;
            }
        }

        double global_I = 0.0;
        MPI_Allreduce(&local_I, &global_I, 1, MPI_DOUBLE, MPI_SUM, fem.H1_fes->GetComm());

        return global_I;
    }
}