module;
#include <algorithm>
#include <optional>
#include <stdexcept>
#include <utility>

export module mean_field:quadrature.policy;
export import :utils.misc;

export namespace mean_field::quadrature {
    enum class Term {
        gravity_hdiv_mass,
        gravity_divergence,
        gravity_source,
        gravity_boundary,
        density_projection,
        mass_conservation,
        center_of_mass,
        quadrupole,
        gravitational_energy,
        virial,
        error_norm
    };

    enum class QuadratureRole {
        discretization,
        preconditioner,
        diagnostic,
        projection
    };

    enum class MappingKind {
        none,
        affine,
        general,
        kelvin
    };

    enum class Mode {
        fast,
        production,
        reference,
        convergence
    };

    struct RuleControl {
        std::optional<int> fixed_order;
        int boost = 0;
    };

    struct RoleControls {
        RuleControl discretization;
        RuleControl preconditioner;
        RuleControl diagnostic;
        RuleControl projection;
    };


    struct RuleSet {
        RuleControl gravity_hdiv_mass;
        RuleControl gravity_divergence;
        RuleControl gravity_source;
        RuleControl gravity_boundary;
        RuleControl density_projection;
        RuleControl mass_conservation;
        RuleControl center_of_mass;
        RuleControl quadrupole;
        RuleControl gravitational_energy;
        RuleControl virial;
        RuleControl error_norm;
        RoleControls roles;
        RuleControl fallback;
    };

    struct Query {
        Term term;
        QuadratureRole role = QuadratureRole::discretization;
        utils::DOMAINS domain = utils::DOMAINS::ALL;
        MappingKind mapping = MappingKind::none;
        int trial_order = 0;
        int test_order = 0;
        int coefficient_order = 0;
        int geometry_weight_order = 0;
        std::optional<int> base_order;
    };

    struct Resolution {
        int base_order;
        int boost;
        int order;
        bool used_fixed_order;
    };
    struct QuadratureTermOptions {
        std::optional<int> fixed_order;
        int additional_boost = 0;
    };

    struct QuadratureManifestOptions {
        bool enabled = false;
        bool include_repeated_queries = false;
        std::optional<std::string> output_file;
    };

    struct QuadratureValidationOptions {
        bool require_explicit_base_order = false;
        bool require_explicit_mfem_rule = false;
        bool reject_negative_boosts = true;
        bool report_unused_overrides = true;
    };

    struct QuadratureRoleOptions {
        QuadratureTermOptions discretization;
        QuadratureTermOptions preconditioner;
        QuadratureTermOptions diagnostic;
        QuadratureTermOptions projection;
    };

    struct QuadratureOptions {
        Mode mode = Mode::production;
        int global_boost = 0;
        std::optional<int> fallback_fixed_order;

        QuadratureTermOptions gravity_hdiv_mass;
        QuadratureTermOptions gravity_divergence;
        QuadratureTermOptions gravity_source;
        QuadratureTermOptions gravity_boundary;
        QuadratureTermOptions density_projection;
        QuadratureTermOptions mass_conservation;
        QuadratureTermOptions center_of_mass;
        QuadratureTermOptions quadrupole;
        QuadratureTermOptions gravitational_energy;
        QuadratureTermOptions virial;
        QuadratureTermOptions error_norm;

        QuadratureRoleOptions roles;

        std::vector<int> convergence_boosts = {0, 2, 4};
        QuadratureManifestOptions manifest;
        QuadratureValidationOptions validation;
    };

    RuleSet make_rule_set(Mode mode, int global_boost = 0);

    class Policy {
    public:
        explicit Policy(RuleSet rule_set);
        Resolution resolve(const Query& query) const;

    private:
        const RuleControl& get_control(Term term) const;
        static int compute_base_order(const Query& query) ;
        const RuleControl& get_role_control(QuadratureRole role) const;

        RuleSet rule_set;
    };

    RuleSet make_rule_set(const Mode mode, const int global_boost) {
        RuleSet rule_set;

        switch (mode) {
            case Mode::fast:
            case Mode::production:
            case Mode::convergence:
                rule_set.fallback.boost = global_boost;
                break;
            case Mode::reference:
                rule_set.fallback.boost = global_boost + 8;
                break;
        }

        return rule_set;
    }

    Policy::Policy(RuleSet rule_set) : rule_set(std::move(rule_set)) {}

    Resolution Policy::resolve(const Query& query) const {
        const int base_order = compute_base_order(query);
        const RuleControl& term_control = get_control(query.term);
        const RuleControl& role_control = get_role_control(query.role);
        std::optional<int> fixed_order;

        if (term_control.fixed_order.has_value()) {
            fixed_order = term_control.fixed_order;
        } else if (role_control.fixed_order.has_value()) {
            fixed_order = role_control.fixed_order;
        } else {
            fixed_order = rule_set.fallback.fixed_order;
        }

        if (fixed_order.has_value()) {
            if (*fixed_order < 0) {
                throw std::invalid_argument("Quadrature fixed order cannot be negative.");
            }

            return {.base_order = base_order, .boost = 0, .order = *fixed_order, .used_fixed_order = true};
        }

        const int boost = rule_set.fallback.boost + role_control.boost + term_control.boost;
        const int order = base_order + boost;

        if (order < 0) {
            throw std::invalid_argument("Resolved quadrature order cannot be negative.");
        }

        return {.base_order = base_order, .boost = boost, .order = order, .used_fixed_order = false};
    }
    const RuleControl& Policy::get_control(const Term term) const {
        switch (term) {
            case Term::gravity_hdiv_mass: return rule_set.gravity_hdiv_mass;
            case Term::gravity_divergence: return rule_set.gravity_divergence;
            case Term::gravity_source: return rule_set.gravity_source;
            case Term::gravity_boundary: return rule_set.gravity_boundary;
            case Term::density_projection: return rule_set.density_projection;
            case Term::mass_conservation: return rule_set.mass_conservation;
            case Term::center_of_mass: return rule_set.center_of_mass;
            case Term::quadrupole: return rule_set.quadrupole;
            case Term::gravitational_energy: return rule_set.gravitational_energy;
            case Term::virial: return rule_set.virial;
            case Term::error_norm: return rule_set.error_norm;
        }

        throw std::logic_error("Unknown quadrature term.");
    }

    int Policy::compute_base_order(const Query& query) {
        if (query.base_order.has_value()) {
            if (*query.base_order < 0) {
                throw std::invalid_argument("Quadrature base order cannot be negative.");
            }
            return *query.base_order;
        }

        if (query.trial_order < 0 || query.test_order < 0 || query.coefficient_order < 0 || query.geometry_weight_order < 0) {
            throw std::invalid_argument("Quadrature query orders cannot be negative.");
        }

        int trial_order = query.trial_order;
        if (query.term == Term::gravity_divergence) {
            trial_order = std::max(0, trial_order - 1);
        }

        return trial_order + query.test_order + query.coefficient_order + query.geometry_weight_order;
    }

    const RuleControl& Policy::get_role_control(const QuadratureRole role) const {
        switch (role) {
            case QuadratureRole::discretization: return rule_set.roles.discretization;
            case QuadratureRole::preconditioner: return rule_set.roles.preconditioner;
            case QuadratureRole::diagnostic: return rule_set.roles.diagnostic;
            case QuadratureRole::projection: return rule_set.roles.projection;
        }

        throw std::logic_error("Unknown quadrature role.");
    }

}