module;
#include "mean_field.h"
#include <array>

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;
        using DomainSchema = utils::domain::CoreEnvelopeVacuumDomainSchema;
        mapping::GridFunctionMappingEvaluator mapping_evaluator(
            *fem.domainMapperStateless, *fem.displacement,
            *fem.compactificationCoordinate
        );

        for (int i = 0; i < fem.mesh->GetNE(); i++) {
            if (!DomainSchema::template attribute_belongs_to<utils::domain::Stellar>(
                    fem.mesh->GetAttribute(i)))
                continue;

            mfem::ElementTransformation *T = fem.mesh->GetElementTransformation(i);
            using DensityField             = field::Field<field::Density>;
            const quadrature::Query query  = DensityField::make_query<field::Density::Form::Quadrupole>(
                quadrature::QuadratureRole::diagnostic, T->OrderW(), std::array<int, 1>{2}, utils::DOMAINS::STELLAR,
                quadrature::MappingKind::general
            );
            const mfem::IntegrationRule &ir = *fem.quadratureFactory->get(query, T->GetGeometryType()).integration_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);

                mapping::VolumeMappingContext mapping_context;
                MFEM_VERIFY(
                    mapping_evaluator.EvaluateVolume(*T, ip, mapping_context) ==
                        mapping::MappingStatus::valid,
                    "Moment-of-inertia integration encountered an invalid mapping."
                );
                const mfem::Vector &x_phys = mapping_context.mapping.physical_position;

                const double r_cyl_sq = x_phys(0) * x_phys(0) + x_phys(1) * x_phys(1);
                const double weight   = mapping_context.quadrature.weight;

                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.mesh->GetComm());

        return global_I;
    }
} // namespace mean_field::physics