// Preserve MeanField's independent role/form routing and native RT polynomial-order regressions.
#include <algorithm>
#include <array>
#include <string_view>

#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/bilininteg.hpp"
#include "mfem/fem/fespace.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/backend/mfem/discretization/quadrature/mfem_resolver.hpp"

namespace
{
    namespace quadrature = serif::discretization::quadrature;
    namespace forms = serif::discretization::forms;
    namespace blocks = serif::discretization::blocks;
    using Mass = forms::mass<blocks::displacement::vector, blocks::displacement::vector>;
    using Divergence = forms::divergence<blocks::displacement::vector, blocks::gravity::potential>;
    using PressureIntegral = forms::source<::serif::tests::equations::Enthalpy, blocks::barotropic_constant::mass_normalization>;
    using PressureForce = forms::divergence<blocks::displacement::vector, ::serif::tests::equations::Enthalpy>;
    using EquilibriumConstant = forms::source<blocks::barotropic_constant::mass_normalization, ::serif::tests::equations::Enthalpy>;

    /** @brief The density projection load has one test operand; its source order is supplied dynamically. */
    struct DensityProjectionSource final : forms::form
    {
        static constexpr std::string_view name = "density projection source";
        using trial_type = blocks::density::mass;
        using test_type = blocks::density::mass;
        using operand_types = serif::utils::types::TypeList<forms::Operand<blocks::density::mass>>;
    };
}

TEST_CASE("MeanField typed quadrature queries preserve source orders and physical metadata", "[unit][quadrature][meanfield]")
{
    namespace domain = serif::discretization::domain;
    const quadrature::Policy policy{quadrature::RuleSet{}};
    const quadrature::backend::mfem::MFEMResolver resolver{policy};
    constexpr int density_order = quadrature::backend::mfem::resolve_operand_order<forms::Operand<blocks::density::mass>>();
    const auto density = quadrature::make_query<DensityProjectionSource>(quadrature::QuadratureRole::projection, domain::StellarDomains{}, quadrature::MappingKind::general, 0, density_order, 4, 3);
    CHECK(density.form == forms::form_descriptor<DensityProjectionSource>);
    CHECK(density.role == quadrature::QuadratureRole::projection);
    CHECK(std::ranges::equal(density.domain.domains(), std::array{domain::CoreDomain::name, domain::EnvelopeDomain::name}));
    CHECK(density.mapping == quadrature::MappingKind::general);
    CHECK(density.geometry_weight_order == 3);
    CHECK(density.trial_order == 0);
    CHECK(density.test_order == 3);
    CHECK(density.coefficient_order == 4);
    CHECK(policy.resolve(density).base_order == 10);
    const std::array dynamic_orders{4};
    const auto density_rule = resolver.resolve<DensityProjectionSource>({.role = density.role, .geometry = mfem::Geometry::CUBE, .geometry_weight_order = density.geometry_weight_order, .dynamic_orders = dynamic_orders, .domain = density.domain, .mapping = density.mapping});
    CHECK(density_rule.resolution.base_order == 10);
    CHECK(density_rule.integration_rule == &mfem::IntRules.Get(mfem::Geometry::CUBE, 10));

    constexpr int constant_order = quadrature::backend::mfem::resolve_operand_order<forms::Operand<blocks::barotropic_constant::mass_normalization>>();
    constexpr int enthalpy_order = quadrature::backend::mfem::resolve_operand_order<forms::Operand<::serif::tests::equations::Enthalpy>>();
    const auto equilibrium = quadrature::make_query<EquilibriumConstant>(quadrature::QuadratureRole::discretization, domain::StellarDomains{}, quadrature::MappingKind::none, constant_order, enthalpy_order, 0, 2);
    CHECK(equilibrium.form == forms::form_descriptor<EquilibriumConstant>);
    CHECK(equilibrium.role == quadrature::QuadratureRole::discretization);
    CHECK(std::ranges::equal(equilibrium.domain.domains(), std::array{domain::CoreDomain::name, domain::EnvelopeDomain::name}));
    CHECK(equilibrium.mapping == quadrature::MappingKind::none);
    CHECK(equilibrium.geometry_weight_order == 2);
    CHECK(equilibrium.trial_order == 0);
    CHECK(equilibrium.test_order == 3);
    CHECK(policy.resolve(equilibrium).base_order == 5);
    const auto equilibrium_rule = resolver.resolve<EquilibriumConstant>({.geometry = mfem::Geometry::CUBE, .geometry_weight_order = 2, .domain = equilibrium.domain});
    CHECK(equilibrium_rule.resolution.base_order == 5);
    CHECK(equilibrium_rule.integration_rule == &mfem::IntRules.Get(mfem::Geometry::CUBE, 5));

    CHECK_THROWS_AS(resolver.resolve<DensityProjectionSource>({.geometry = mfem::Geometry::CUBE, .geometry_weight_order = -1}), std::invalid_argument);
    const std::array invalid_orders{-1};
    CHECK_THROWS_AS(resolver.resolve<PressureForce>({.geometry = mfem::Geometry::CUBE, .dynamic_orders = invalid_orders}), std::invalid_argument);
}

TEST_CASE("MeanField pressure integral and pressure force retain independent quadrature controls", "[unit][quadrature][meanfield]")
{
    quadrature::RuleSet rules;
    rules.set_form_control<PressureIntegral>({.boost = 3});
    rules.set_form_control<PressureForce>({.boost = 5});
    const quadrature::backend::mfem::MFEMResolver resolver{quadrature::Policy{rules}};
    const std::array dynamic_orders{9};
    const auto integral = resolver.resolve<PressureIntegral>({.role = quadrature::QuadratureRole::diagnostic, .geometry = mfem::Geometry::CUBE, .geometry_weight_order = 2, .dynamic_orders = dynamic_orders, .domain = serif::discretization::domain::StellarDomains{}, .mapping = quadrature::MappingKind::general});
    const auto force = resolver.resolve<PressureForce>({.geometry = mfem::Geometry::CUBE, .geometry_weight_order = 2, .dynamic_orders = dynamic_orders, .domain = serif::discretization::domain::StellarDomains{}, .mapping = quadrature::MappingKind::general});
    CHECK(integral.resolution.base_order == 14);
    CHECK(integral.resolution.boost == 3);
    CHECK(integral.resolution.order == 17);
    CHECK(force.resolution.base_order == 16);
    CHECK(force.resolution.boost == 5);
    CHECK(force.resolution.order == 21);
    CHECK(integral.integration_rule == &mfem::IntRules.Get(mfem::Geometry::CUBE, 17));
    CHECK(force.integration_rule == &mfem::IntRules.Get(mfem::Geometry::CUBE, 21));
}

TEST_CASE("MeanField quadrature routing generalizes to independent typed forms and every numerical role", "[unit][quadrature][meanfield]")
{
    const std::array descriptors{forms::form_descriptor<Mass>, forms::form_descriptor<Divergence>, forms::form_descriptor<PressureIntegral>, forms::form_descriptor<PressureForce>, forms::form_descriptor<forms::mass<blocks::density::mass, blocks::density::mass>>, forms::form_descriptor<forms::source<blocks::density::mass, blocks::gravity::potential>>, forms::form_descriptor<forms::source<blocks::gravity::potential, blocks::density::mass>>, forms::form_descriptor<forms::projection<::serif::tests::equations::Enthalpy, blocks::density::mass>>};
    const std::array roles{quadrature::QuadratureRole::discretization, quadrature::QuadratureRole::preconditioner, quadrature::QuadratureRole::diagnostic, quadrature::QuadratureRole::projection};
    quadrature::RuleSet rules;
    rules.set_fallback_control({.boost = 3});
    for (std::size_t form = 0; form < descriptors.size(); ++form)
    {
        rules.set_form_control(descriptors[form], {.boost = static_cast<int>(form + 1)});
    }
    for (std::size_t role = 0; role < roles.size(); ++role)
    {
        rules.set_role_control(roles[role], {.boost = static_cast<int>(2 * role)});
    }
    const quadrature::Policy policy{rules};
    for (std::size_t form = 0; form < descriptors.size(); ++form)
    {
        for (std::size_t role = 0; role < roles.size(); ++role)
        {
            const auto result = policy.resolve({.form = descriptors[form], .role = roles[role], .base_order = 20});
            CHECK(result.base_order == 20);
            CHECK(result.boost == static_cast<int>(3 + form + 1 + 2 * role));
            CHECK(result.order == 20 + result.boost);
            CHECK_FALSE(result.used_fixed_order);
        }
    }
    // Replacing an existing rule must not leave an earlier rule shadowing the new physical choice.
    rules.set_form_control<PressureForce>({.fixed_order = 31});
    CHECK(quadrature::Policy{rules}.resolve({.form = forms::form_descriptor<PressureForce>, .base_order = 20}).order == 31);
    CHECK(quadrature::Policy{rules}.resolve({.form = forms::form_descriptor<PressureIntegral>, .base_order = 20}).order == 26);
}