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

export module mean_field:quadrature.mfem;
export import :quadrature.policy;
export import :integrators.centrifugal;
import :field.mfem;

export namespace mean_field::quadrature {
    struct MfemRule {
        Resolution resolution;
        const mfem::IntegrationRule *integration_rule;
    };

    class RuleFactory {
    public:
        explicit RuleFactory(Policy policy);
        MfemRule
        get(const Query &query,
            mfem::Geometry::Type geometry) const;
        MfemRule
        get(Term term,
            QuadratureRole role,
            mfem::Geometry::Type geometry,
            int base_order,
            utils::DOMAINS domain = utils::DOMAINS::ALL,
            MappingKind mapping   = MappingKind::none) const;

        Resolution configure_gravity_hdiv_mass(
            mfem::VectorFEMassIntegrator &integrator,
            QuadratureRole role,
            const mfem::FiniteElement &element,
            const mfem::ElementTransformation &transformation,
            utils::DOMAINS domain = utils::DOMAINS::ALL,
            MappingKind mapping   = MappingKind::none
        ) const;

        Resolution configure_gravity_divergence(
            mfem::VectorFEDivergenceIntegrator &integrator,
            QuadratureRole role,
            const mfem::FiniteElement &trial_element,
            const mfem::FiniteElement &test_element,
            const mfem::ElementTransformation &transformation,
            utils::DOMAINS domain = utils::DOMAINS::ALL,
            MappingKind mapping   = MappingKind::none
        ) const;

        Resolution configure_gravity_boundary(
            mfem::VectorFEBoundaryFluxLFIntegrator &integrator,
            QuadratureRole role,
            const mfem::FiniteElement &boundary_element,
            utils::DOMAINS domain = utils::DOMAINS::VACUUM,
            MappingKind mapping   = MappingKind::none
        ) const;

        Resolution configure_gravity_source(
            mfem::DomainLFIntegrator &integrator,
            QuadratureRole role,
            const mfem::FiniteElement &test_element,
            const mfem::ElementTransformation &transformation,
            int coefficient_order,
            utils::DOMAINS domain = utils::DOMAINS::STELLAR,
            MappingKind mapping   = MappingKind::none
        ) const;

        Resolution configure_gravity_source(
            mfem::MixedScalarMassIntegrator &integrator,
            QuadratureRole role,
            const mfem::FiniteElement &trial_element,
            const mfem::FiniteElement &test_element,
            const mfem::ElementTransformation &transformation,
            int coefficient_order,
            utils::DOMAINS domain = utils::DOMAINS::STELLAR,
            MappingKind mapping   = MappingKind::none
        ) const;

        Resolution configure_centrifugal(
            integrators::CentrifugalForceIntegrator &integrator,
            QuadratureRole role,
            const mfem::FiniteElement &density_element,
            const mfem::FiniteElement &velocity_element,
            const mfem::ElementTransformation &transformation,
            int position_order,
            utils::DOMAINS domain = utils::DOMAINS::STELLAR,
            MappingKind mapping   = MappingKind::none
        ) const;

        template <typename IntegratorType>
        Resolution configure(
            IntegratorType &integrator,
            Term term,
            QuadratureRole role,
            mfem::Geometry::Type geometry,
            int base_order,
            utils::DOMAINS domain = utils::DOMAINS::ALL,
            MappingKind mapping   = MappingKind::none
        ) const;

    private:
        Policy policy;
    };

    RuleFactory::RuleFactory(Policy policy) : policy(std::move(policy)) {
    }

    MfemRule RuleFactory::get(
        const Query &query,
        const mfem::Geometry::Type geometry
    ) const {
        const Resolution resolution                   = policy.resolve(query);
        const mfem::IntegrationRule &integration_rule = mfem::IntRules.Get(geometry, resolution.order);
        return {.resolution = resolution, .integration_rule = &integration_rule};
    }

    MfemRule RuleFactory::get(
        const Term term,
        const QuadratureRole role,
        const mfem::Geometry::Type geometry,
        const int base_order,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        Query query{.term = term};
        query.domain     = domain;
        query.mapping    = mapping;
        query.role       = role;
        query.base_order = base_order;
        return get(query, geometry);
    }

    Resolution RuleFactory::configure_gravity_hdiv_mass(
        mfem::VectorFEMassIntegrator &integrator,
        const QuadratureRole role,
        const mfem::FiniteElement &element,
        const mfem::ElementTransformation &transformation,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        using GravityField = field::Field<field::Gravity>;
        MFEM_VERIFY(
            element.GetOrder() == field::Gravity::Flux::familyOrder + 1,
            "The H(div) element order does not match the registered gravity "
            "flux."
        );
        const Query query = GravityField::make_query<field::Gravity::Form::HDivMass>(
            role, transformation.OrderW(), {}, domain, mapping
        );
        const auto [resolution, integration_rule] = get(query, element.GetGeomType());
        integrator.SetIntegrationRule(*integration_rule);
        return resolution;
    }

    Resolution RuleFactory::configure_gravity_divergence(
        mfem::VectorFEDivergenceIntegrator &integrator,
        const QuadratureRole role,
        const mfem::FiniteElement &trial_element,
        const mfem::FiniteElement &test_element,
        const mfem::ElementTransformation &transformation,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        using GravityField = field::Field<field::Gravity>;
        MFEM_VERIFY(
            trial_element.GetOrder() == field::Gravity::Flux::familyOrder + 1,
            "The divergence trial element does not match the registered "
            "gravity flux."
        );
        MFEM_VERIFY(
            test_element.GetOrder() == field::Gravity::Potential::familyOrder,
            "The divergence test element does not match the registered "
            "gravity potential."
        );
        const Query query = GravityField::make_query<field::Gravity::Form::DivergenceCoupling>(
            role, transformation.OrderW(), {}, domain, mapping
        );

        const auto [resolution, integration_rule] = get(query, trial_element.GetGeomType());
        integrator.SetIntegrationRule(*integration_rule);
        return resolution;
    }

    Resolution RuleFactory::configure_gravity_boundary(
        mfem::VectorFEBoundaryFluxLFIntegrator &integrator,
        const QuadratureRole role,
        const mfem::FiniteElement &boundary_element,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        using GravityField = field::Field<field::Gravity>;
        MFEM_VERIFY(
            boundary_element.GetOrder() == field::Gravity::Flux::familyOrder,
            "The boundary element does not match the registered gravity-flux "
            "normal trace."
        );
        const Query query = GravityField::make_query<field::Gravity::Form::Boundary>(role, 0, {}, domain, mapping);
        const auto [resolution, integration_rule] = get(query, boundary_element.GetGeomType());
        integrator.SetIntegrationRule(*integration_rule);
        return resolution;
    }

    Resolution RuleFactory::configure_gravity_source(
        mfem::DomainLFIntegrator &integrator,
        const QuadratureRole role,
        const mfem::FiniteElement &test_element,
        const mfem::ElementTransformation &transformation,
        const int coefficient_order,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        using GravityField = field::Field<field::Gravity>;
        MFEM_VERIFY(
            test_element.GetOrder() == field::Gravity::Potential::familyOrder,
            "The gravity-source test element does not match the registered "
            "gravity potential."
        );
        MFEM_VERIFY(
            coefficient_order == field::Density::Scalar::familyOrder,
            "The gravity-source coefficient order does not match the "
            "registered density field."
        );
        const Query query = GravityField::make_query<field::Gravity::Form::SourceLinear>(
            role, transformation.OrderW(), {}, domain, mapping
        );

        const auto [resolution, integration_rule] = get(query, test_element.GetGeomType());
        integrator.SetIntegrationRule(*integration_rule);
        return resolution;
    }
    Resolution RuleFactory::configure_gravity_source(
        mfem::MixedScalarMassIntegrator &integrator,
        QuadratureRole role,
        const mfem::FiniteElement &trial_element,
        const mfem::FiniteElement &test_element,
        const mfem::ElementTransformation &transformation,
        int coefficient_order,
        utils::DOMAINS domain,
        MappingKind mapping
    ) const {
        MFEM_VERIFY(
            trial_element.GetGeomType() == test_element.GetGeomType(),
            "Gravity source trial and test elements must use the same geometry."
        );
        MFEM_VERIFY(
            trial_element.GetGeomType() == transformation.GetGeometryType(),
            "Gravity source element and transformation geometries must agree."
        );

        using GravityField = field::Field<field::Gravity>;
        MFEM_VERIFY(
            trial_element.GetOrder() == field::Density::Scalar::familyOrder,
            "The gravity-source trial element does not match the registered "
            "density field."
        );
        MFEM_VERIFY(
            test_element.GetOrder() == field::Gravity::Potential::familyOrder,
            "The gravity-source test element does not match the registered "
            "gravity potential."
        );
        MFEM_VERIFY(
            coefficient_order == 0, "The mapped gravity-source coefficient order must be zero; "
                                    "density order is supplied by the registered trial field."
        );
        const Query query = GravityField::make_query<field::Gravity::Form::SourceProjection>(
            role, transformation.OrderW(), {}, domain, mapping
        );

        const auto [resolution, integration_rule] = get(query, transformation.GetGeometryType());
        integrator.SetIntRule(integration_rule);
        return resolution;
    }

    Resolution RuleFactory::configure_centrifugal(
        integrators::CentrifugalForceIntegrator &integrator,
        const QuadratureRole role,
        const mfem::FiniteElement &density_element,
        const mfem::FiniteElement &velocity_element,
        const mfem::ElementTransformation &transformation,
        const int position_order,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        const Query query = {
            .term                  = Term::centrifugal,
            .role                  = role,
            .domain                = domain,
            .mapping               = mapping,
            .trial_order           = density_element.GetOrder(),
            .test_order            = velocity_element.GetOrder(),
            .coefficient_order     = position_order,
            .geometry_weight_order = transformation.OrderW()
        };

        const auto [resolution, integration_rule] = get(query, velocity_element.GetGeomType());
        integrator.SetIntegrationRule(*integration_rule);
        return resolution;
    }

    template <typename IntegratorType>
    Resolution RuleFactory::configure(
        IntegratorType &integrator,
        const Term term,
        const QuadratureRole role,
        const mfem::Geometry::Type geometry,
        const int base_order,
        const utils::DOMAINS domain,
        const MappingKind mapping
    ) const {
        const auto [resolution, integration_rule] = get(term, role, geometry, base_order, domain, mapping);
        integrator.SetIntegrationRule(*integration_rule);
        return resolution;
    }

} // namespace mean_field::quadrature