module;
#include <array>
#include <mfem.hpp>

module mean_field;
import :mapping.coefficients;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

    mfem::Array<int> make_domain_marker(
        const mfem::Mesh &mesh,
        const mean_field::utils::DOMAINS domain
    ) {
        switch (domain) {
        case mean_field::utils::DOMAINS::CORE:
            return mean_field::utils::domain::make_attribute_marker<mean_field::utils::domain::Core, DomainSchema>(
                mesh
            );
        case mean_field::utils::DOMAINS::ENVELOPE:
            return mean_field::utils::domain::make_attribute_marker<mean_field::utils::domain::Envelope, DomainSchema>(
                mesh
            );
        case mean_field::utils::DOMAINS::ALL:
            return mean_field::utils::domain::make_attribute_marker<mean_field::utils::domain::All, DomainSchema>(mesh);
        case mean_field::utils::DOMAINS::STELLAR:
            return mean_field::utils::domain::make_attribute_marker<mean_field::utils::domain::Stellar, DomainSchema>(
                mesh
            );
        case mean_field::utils::DOMAINS::VACUUM:
            return mean_field::utils::domain::make_attribute_marker<mean_field::utils::domain::Vacuum, DomainSchema>(
                mesh
            );
        }
        MFEM_ABORT("Unsupported integration domain.");
    }

    template <typename FormT>
    const mfem::IntegrationRule &get_density_rule(
        const mean_field::fem::FEM &fem,
        const mfem::ElementTransformation &transformation,
        const std::array<
            int,
            FormT::dynamicOrderCount> &dynamic_orders = {},
        const mean_field::utils::DOMAINS domain       = mean_field::utils::DOMAINS::ALL
    ) {
        using DensityField                        = mean_field::field::Field<mean_field::field::Density>;

        const mean_field::quadrature::Query query = DensityField::make_query<FormT>(
            mean_field::quadrature::QuadratureRole::diagnostic, transformation.OrderW(), dynamic_orders, domain,
            fem.has_mapping() ? mean_field::quadrature::MappingKind::general : mean_field::quadrature::MappingKind::none
        );

        return *fem.quadratureFactory->get(query, transformation.GetGeometryType()).integration_rule;
    }
} // namespace

namespace mean_field::analysis {
    double domain_integrate_grid_function(
        const fem::FEM &fem,
        const mfem::GridFunction &gf,
        utils::DOMAINS domain,
        mapping::COORDINATE_SPACE coord_space
    ) {
        mfem::LinearForm lf(fem.densityFes.get());
        mfem::GridFunctionCoefficient gf_c(&gf);
        double local_integral;
        mfem::Array<int> elem_markers                                    = make_domain_marker(*fem.mesh, domain);
        const mfem::ElementTransformation &representative_transformation = *fem.mesh->GetElementTransformation(0);
        const mfem::IntegrationRule &integration_rule =
            get_density_rule<field::Density::Form::MassConservation>(fem, representative_transformation, {}, domain);

        if (fem.has_mapping() && coord_space == mapping::COORDINATE_SPACE::PHYSICAL) {
            mapping::MappedScalarCoefficient mapped_gf_c(
                *fem.domainMapperStateless, *fem.displacement, *fem.compactificationCoordinate, gf_c
            );

            // ReSharper disable once CppDFAMemoryLeak // Disabled because MFEM
            // takes ownership so memory is not leaked
            auto *lf_integrator = new mfem::DomainLFIntegrator(mapped_gf_c);
            lf_integrator->SetIntRule(&integration_rule);
            lf.AddDomainIntegrator(lf_integrator, elem_markers);
            lf.Assemble();

            local_integral = lf.Sum();
        } else {
            if (coord_space == mapping::COORDINATE_SPACE::PHYSICAL) {
                MFEM_ABORT(
                    "Physical evaluation mode requested but no mapping "
                    "provided. Check "
                    "domain bounds and mapping setup."
                );
            }
            auto *lf_integrator = new mfem::DomainLFIntegrator(gf_c);
            lf_integrator->SetIntRule(&integration_rule);
            lf.AddDomainIntegrator(lf_integrator, elem_markers);
            lf.Assemble();
            local_integral = lf.Sum();
        }

        double global_integral = 0.0;
        MPI_Allreduce(&local_integral, &global_integral, 1, MPI_DOUBLE, MPI_SUM, fem.mesh->GetComm());
        return global_integral;
    }

    mfem::Vector get_com(
        const fem::FEM &fem,
        const mfem::GridFunction &rho
    ) {
        const int dim = fem.mesh->Dimension();
        mapping::GridFunctionMappingEvaluator mapping_evaluator(
            *fem.domainMapperStateless, *fem.displacement, *fem.compactificationCoordinate
        );
        mfem::Vector local_com(dim);
        local_com         = 0.0;
        double local_mass = 0.0;

        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 *trans = fem.mesh->GetElementTransformation(i);
            const mfem::IntegrationRule &ir    = get_density_rule<field::Density::Form::CenterOfMass>(
                fem, *trans, std::array<int, 1>{1}, utils::DOMAINS::STELLAR
            );

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

                mapping::VolumeMappingContext mapping_context;
                MFEM_VERIFY(
                    mapping_evaluator.EvaluateVolume(*trans, ip, mapping_context) == mapping::MappingStatus::valid,
                    "Center-of-mass integration encountered an invalid mapping."
                );
                const double weight            = mapping_context.quadrature.weight;
                double rho_val                 = rho.GetValue(i, ip);

                const mfem::Vector &phys_point = mapping_context.mapping.physical_position;

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

                for (int d = 0; d < dim; ++d) {
                    local_com(d) += phys_point(d) * mass_term;
                }
            }
        }

        double global_mass = 0.0;
        mfem::Vector global_com(dim);
        MPI_Comm comm = fem.mesh->GetComm();

        MPI_Allreduce(&local_mass, &global_mass, 1, MPI_DOUBLE, MPI_SUM, comm);

        MPI_Allreduce(local_com.GetData(), global_com.GetData(), dim, MPI_DOUBLE, MPI_SUM, comm);

        if (global_mass > 1e-18) {
            global_com /= global_mass;
        } else {
            global_com = 0.0;
        }

        return global_com;
    }

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

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

        std::unique_ptr<mfem::Coefficient> s2_coeff;
        if (fem.has_mapping()) {
            s2_coeff = std::make_unique<mapping::PhysicalPositionFunctionCoefficient>(
                *fem.domainMapperStateless, *fem.displacement, *fem.compactificationCoordinate, s2_func
            );
        } else {
            s2_coeff = std::make_unique<mfem::FunctionCoefficient>(s2_func);
        }

        mfem::GridFunctionCoefficient rho_coeff(&rho);
        mfem::ProductCoefficient I_integrand(rho_coeff, *s2_coeff);

        mfem::LinearForm I_lf(fem.densityFes.get());
        const mfem::ElementTransformation &representative_transformation = *fem.mesh->GetElementTransformation(0);
        const mfem::IntegrationRule &integration_rule = get_density_rule<field::Density::Form::Quadrupole>(
            fem, representative_transformation, std::array<int, 1>{2}, utils::DOMAINS::STELLAR
        );
        mfem::Array<int> stellar_markers =
            utils::domain::make_attribute_marker<utils::domain::Stellar, DomainSchema>(*fem.mesh);

        double local_I = 0.0;
        if (fem.has_mapping()) {
            mapping::MappedScalarCoefficient mapped_integrand(
                *fem.domainMapperStateless, *fem.displacement, *fem.compactificationCoordinate, I_integrand
            );
            auto *integrator = new mfem::DomainLFIntegrator(mapped_integrand);
            integrator->SetIntRule(&integration_rule);
            I_lf.AddDomainIntegrator(integrator, stellar_markers);
            I_lf.Assemble();
            local_I = I_lf.Sum();
        } else {
            auto *integrator = new mfem::DomainLFIntegrator(I_integrand);
            integrator->SetIntRule(&integration_rule);
            I_lf.AddDomainIntegrator(integrator, stellar_markers);
            I_lf.Assemble();
            local_I = I_lf.Sum();
        }

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

    double get_mesh_volume(
        const fem::FEM &fem,
        const mapping::COORDINATE_SPACE coordinate_space,
        const utils::DOMAINS domain
    ) {
        mfem::ParMesh &mesh = *fem.mesh;
        const bool physical = (coordinate_space == mapping::COORDINATE_SPACE::PHYSICAL);

        if (physical && !fem.has_mapping()) {
            MFEM_ABORT("Physical volume requested but no domain mapping is available.");
        }

        double local_volume = 0.0;
        mapping::GridFunctionMappingEvaluator mapping_evaluator(
            *fem.domainMapperStateless, *fem.displacement, *fem.compactificationCoordinate
        );

        for (int e = 0; e < mesh.GetNE(); ++e) {
            const int attr      = mesh.GetAttribute(e);
            const bool selected = domain == utils::DOMAINS::ALL ||
                                  (domain == utils::DOMAINS::STELLAR &&
                                   DomainSchema::template attribute_belongs_to<utils::domain::Stellar>(attr)) ||
                                  (domain == utils::DOMAINS::VACUUM &&
                                   DomainSchema::template attribute_belongs_to<utils::domain::Vacuum>(attr));
            if (!selected)
                continue;
            mfem::ElementTransformation *T = mesh.GetElementTransformation(e);
            const mfem::IntegrationRule &ir =
                get_density_rule<field::Density::Form::MassConservation>(fem, *T, {}, domain);

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

                double dV = ip.weight * T->Weight();

                if (physical) {
                    mapping::VolumeMappingContext context;
                    MFEM_VERIFY(
                        mapping_evaluator.EvaluateVolume(*T, ip, context) == mapping::MappingStatus::valid,
                        "Mesh-volume integration encountered an invalid mapping."
                    );
                    dV = context.quadrature.weight;
                }

                local_volume += dV;
            }
        }

        double global_volume = 0.0;
        MPI_Allreduce(&local_volume, &global_volume, 1, MPI_DOUBLE, MPI_SUM, mesh.GetComm());
        return global_volume;
    }
} // namespace mean_field::analysis