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

export module mean_field:quadrature.mfem;
export import :quadrature.policy;

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;

        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 {
        const int base_order = 2 * element.GetOrder() + transformation.OrderW();
        return configure(integrator, Term::gravity_hdiv_mass, role, element.GetGeomType(), base_order, domain, mapping);
    }

    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 {
        const Query query = {
            .term = Term::gravity_divergence,
            .role = role,
            .domain = domain,
            .mapping = mapping,
            .trial_order = trial_element.GetOrder(),
            .test_order = test_element.GetOrder(),
            .geometry_weight_order = transformation.OrderW()
        };

        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 {
        const int base_order = 2 * boundary_element.GetOrder();
        return configure(integrator, Term::gravity_boundary, role, boundary_element.GetGeomType(), base_order, domain, mapping);
    }

    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 {
        const Query query = {
            .term = Term::gravity_source,
            .role = role,
            .domain = domain,
            .mapping = mapping,
            .test_order = test_element.GetOrder(),
            .coefficient_order = coefficient_order,
            .geometry_weight_order = transformation.OrderW()
        };

        const auto [resolution, integration_rule] = get(query, test_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;
    }

}