module;
#include <algorithm>
#include <array>
#include <catch2/internal/catch_stringref.hpp>
#include <concepts>
#include <memory>
#include <string>

#include <mfem.hpp>
#include <optional>
#include <utility>
export module test_helpers;
import mean_field;

template <std::size_t N> struct Tag {
    std::array<char, N> chars{};

    // ReSharper disable once CppNonExplicitConvertingConstructor
    consteval Tag(
        std::array<
            char,
            N> arr
    )
        : chars(arr) {
    }

    // ReSharper disable once CppNonExplicitConversionOperator
    constexpr operator const char *() const {
        return chars.data();
    }

    // ReSharper disable once CppNonExplicitConversionOperator
    constexpr operator Catch::StringRef() const {
        return Catch::StringRef(chars.data(), N - 1);
    }

    template <std::size_t M> consteval Tag<N + M - 1> operator&(const Tag<M> &other) const {
        std::array<char, N + M - 1> res{};
        std::ranges::copy(chars.begin(), chars.end() - 1, res.begin());
        std::ranges::copy(other.chars, res.begin() + (N - 1));
        return {res};
    }
};

template <std::size_t N> consteval auto make_tag(const char (&str)[N]) {
    std::array<char, N + 2> res{};
    res[0] = '[';
    std::ranges::copy(str, str + N - 1, res.begin() + 1);
    res[N]     = ']';
    res[N + 1] = '\0';
    return Tag<N + 2>{res};
}

template <
    std::size_t N,
    std::size_t M>
consteval auto sub_tag(
    const Tag<N> &parent,
    const char (&str)[M]
) {
    return parent & make_tag(str);
}

namespace test_utils::detail {
    std::optional<mean_field::utils::Args> configured_args;

    mean_field::utils::Args make_default_args() {
        mean_field::utils::Args args;
        args.mesh_file = "sandbox.smesh";
        args.p.rtol    = 1.0e-12;
        args.p.atol    = 1.0e-12;
        return args;
    }
} // namespace test_utils::detail

export namespace test_utils {
    void set_args(mean_field::utils::Args args) {
        detail::configured_args = std::move(args);
    }

    mean_field::utils::Args setup_args() {
        if (detail::configured_args.has_value()) {
            return *detail::configured_args;
        }

        return detail::make_default_args();
    }
} // namespace test_utils

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

    template <typename FieldT> inline mean_field::field::FieldDofMap make_field_map(const mean_field::fem::FEM &f) {
        if constexpr (std::same_as<FieldT, mean_field::field::Density>) {
            return mean_field::field::make_field_dof_map<FieldT, DomainSchema>(*f.densityFes);
        } else if constexpr (std::same_as<FieldT, mean_field::field::Displacement>) {
            return mean_field::field::make_field_dof_map<FieldT, DomainSchema>(*f.displacementFes);
        } else {
            static_assert(std::same_as<FieldT, mean_field::field::Gravity>);
            return mean_field::field::make_field_dof_map<FieldT, DomainSchema>(*f.gravityFluxFes);
        }
    }

    template <typename FieldT>
    inline mfem::Vector gather_field(
        const mean_field::fem::FEM &f,
        const mfem::Vector &true_vector
    ) {
        return make_field_map<FieldT>(f).gather(true_vector);
    }

    inline mfem::Vector make_deterministic_vector(
        const int size,
        const double phase = 0.0
    ) {
        mfem::Vector vector(size);

        for (int i = 0; i < size; ++i) {
            const double index = static_cast<double>(i + 1);
            vector(i)          = std::sin(0.37 * index + phase) + 0.31 * std::cos(0.19 * index - 0.5 * phase);
        }

        return vector;
    }

    inline mfem::Vector make_displacement(
        const mean_field::fem::FEM &f,
        const double scale
    ) {
        mfem::ParGridFunction displacement(f.displacementFes.get());

        auto displacement_function = [scale](const mfem::Vector &position, mfem::Vector &value) {
            value.SetSize(3);
            value(0) = scale * (0.04 * position(0) + 0.01 * position(1) * position(2));
            value(1) = scale * (-0.03 * position(1) + 0.008 * position(0) * position(2));
            value(2) = scale * (0.02 * position(2) - 0.006 * position(0) * position(1));
        };

        mfem::VectorFunctionCoefficient coefficient(f.mesh->Dimension(), displacement_function);
        displacement.ProjectCoefficient(coefficient);

        mfem::Vector displacement_true;
        displacement.GetTrueDofs(displacement_true);
        return displacement_true;
    }

    inline mfem::Vector make_domain_supported_density(
        const mean_field::fem::FEM &f,
        const bool stellar
    ) {
        mfem::Vector attribute_values(f.mesh->attributes.Max());
        attribute_values   = 0.0;

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

        for (int i = 0; i < f.mesh->attributes.Size(); ++i) {
            const int attribute = f.mesh->attributes[i];
            const bool is_stellar =
                DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Stellar>(attribute);

            if (is_stellar == stellar) {
                attribute_values(attribute - 1) = 1.0;
            }
        }

        mfem::PWConstCoefficient coefficient(attribute_values);
        mfem::ParGridFunction density(f.densityFes.get());
        density.ProjectCoefficient(coefficient);

        mfem::Vector density_true;
        density.GetTrueDofs(density_true);
        return density_true;
    }

    inline mfem::Vector linear_combination(
        const mfem::Vector &first,
        const double first_scale,
        const mfem::Vector &second,
        const double second_scale
    ) {
        MFEM_VERIFY(first.Size() == second.Size(), "Cannot combine vectors with different sizes.");

        mfem::Vector combination(first);
        combination *= first_scale;
        combination.Add(second_scale, second);
        return combination;
    }

    inline double global_norm(
        const mfem::Vector &vector,
        MPI_Comm communicator
    ) {
        const double local_norm_squared = vector * vector;
        double global_norm_squared      = 0.0;
        MPI_Allreduce(&local_norm_squared, &global_norm_squared, 1, MPI_DOUBLE, MPI_SUM, communicator);
        return std::sqrt(global_norm_squared);
    }

    inline double global_dot(
        const mfem::Vector &first,
        const mfem::Vector &second,
        MPI_Comm communicator
    ) {
        MFEM_VERIFY(first.Size() == second.Size(), "Cannot take the dot product of vectors with different sizes.");

        const double local_dot = first * second;
        double global_dot      = 0.0;
        MPI_Allreduce(&local_dot, &global_dot, 1, MPI_DOUBLE, MPI_SUM, communicator);
        return global_dot;
    }

    inline double relative_error(
        const mfem::Vector &computed,
        const mfem::Vector &reference,
        MPI_Comm communicator
    ) {
        MFEM_VERIFY(computed.Size() == reference.Size(), "Cannot compare vectors with different sizes.");

        mfem::Vector difference(computed);
        difference -= reference;

        return global_norm(difference, communicator) /
               std::max(global_norm(reference, communicator), std::numeric_limits<double>::epsilon());
    }

    inline double relative_scalar_error(
        const double computed,
        const double reference
    ) {
        return std::abs(computed - reference) / std::max(std::abs(reference), std::numeric_limits<double>::epsilon());
    }
} // namespace gravity_prepared_test_utils

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

    inline mean_field::mapping::DomainMapper make_domain_mapper() {
        const mean_field::utils::Args args = test_utils::setup_args();
        return mean_field::mapping::DomainMapper(
            args.domain_mapper_options,
            std::make_unique<const mean_field::mapping::compactification::KelvinCompactification>(args.kelvin_options)
        );
    }

    inline constexpr int vacuum_material_attribute =
        DomainSchema::template material_attribute<mean_field::utils::domain::Vacuum>();

    template <typename FieldT>
    inline mean_field::field::FieldDofMap make_map(const mfem::ParFiniteElementSpace &finiteElementSpace) {
        return mean_field::field::make_field_dof_map<FieldT, DomainSchema>(finiteElementSpace);
    }

    template <typename FieldT>
    inline mfem::Vector make_deterministic_supported_vector(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const double phase
    ) {
        const mean_field::field::FieldDofMap map = make_map<FieldT>(finiteElementSpace);
        const mfem::Vector full = gravity_prepared_test_utils::make_deterministic_vector(map.full_size(), phase);
        return map.gather(full);
    }

    inline mfem::Vector make_supported_displacement(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        const mean_field::field::FieldDofMap map = make_map<mean_field::field::Displacement>(*f.displacementFes);
        return map.gather(gravity_prepared_test_utils::make_displacement(f, phase));
    }

    inline void apply_hydrostatic_reference(
        const mean_field::fem::FEM &f,
        const mean_field::physics::RigidRotation &rotation,
        const mfem::Vector &enthalpy,
        const mfem::Vector &gravityPotential,
        const mfem::Vector &displacement,
        const double bernoulliConstant,
        mfem::Vector &residual
    ) {
        const mean_field::field::FieldDofMap enthalpyMap = make_map<mean_field::field::Enthalpy>(*f.enthalpyFes);
        const mean_field::field::FieldDofMap gravityPotentialMap =
            make_map<mean_field::field::Gravity>(*f.gravityPotentialFes);
        const mean_field::field::FieldDofMap displacementMap =
            make_map<mean_field::field::Displacement>(*f.displacementFes);

        mfem::Vector enthalpyTrue(enthalpyMap.full_size());
        mfem::Vector gravityPotentialTrue(gravityPotentialMap.full_size());
        mfem::Vector displacementTrue(displacementMap.full_size());
        mfem::Vector residualTrue;

        enthalpyMap.scatter(enthalpy, enthalpyTrue);
        gravityPotentialMap.scatter(gravityPotential, gravityPotentialTrue);
        displacementMap.scatter(displacement, displacementTrue);

        mean_field::operators::kernels::apply_hydrostatic_equilibrium(
            f, *f.domainMapperStateless, rotation, enthalpyTrue, gravityPotentialTrue, displacementTrue,
            bernoulliConstant, residualTrue
        );

        residual.SetSize(enthalpyMap.reduced_size());
        enthalpyMap.gather(residualTrue, residual);
    }
} // namespace field_dof_test_utils

export namespace tags {
    inline constexpr auto geometry                     = make_tag("geometry");
    inline constexpr auto physics                      = make_tag("physics");
    inline constexpr auto unit                         = make_tag("unit");
    inline constexpr auto mesh                         = make_tag("mesh");
    inline constexpr auto integration                  = make_tag("integration");
    inline constexpr auto solver                       = make_tag("solver");
    inline constexpr auto integrator                   = make_tag("integrator");
    inline constexpr auto mapping                      = make_tag("mapping");
    inline constexpr auto utils                        = make_tag("utils");
    inline constexpr auto mfem_operators               = make_tag("operators");
    inline constexpr auto initialization               = make_tag("initialization");
    inline constexpr auto accuracy                     = make_tag("accuracy");
    inline constexpr auto closure                      = make_tag("closure");
    inline constexpr auto kernels                      = make_tag("kernels");
    inline constexpr auto surface                      = make_tag("surface");
    inline constexpr auto model                        = make_tag("model");

    inline constexpr auto field                        = sub_tag(mesh & physics, "field");
    inline constexpr auto field_dof                    = field & make_tag("dof");
    inline constexpr auto field_dof_unit               = field_dof & unit;
    inline constexpr auto field_dof_integration        = field_dof & integration;

    inline constexpr auto pressure                     = sub_tag(physics, "pressure");

    inline constexpr auto hydro                        = sub_tag(physics, "hydro");
    inline constexpr auto jacobian                     = sub_tag(integration & physics, "jacobian");
    inline constexpr auto residuals                    = sub_tag(integration & physics, "residuals");
    inline constexpr auto volume                       = sub_tag(mesh & geometry, "volume");
    inline constexpr auto quadrature                   = sub_tag(mesh & geometry & solver, "quadrature");
    inline constexpr auto convergence                  = sub_tag(solver, "convergence");
    inline constexpr auto transformations              = sub_tag(mesh & geometry, "transformations");

    inline constexpr auto h_refinement                 = sub_tag(mesh & convergence, "h_refinement");
    inline constexpr auto p_refinement                 = sub_tag(mesh & convergence, "p_refinement");

    inline constexpr auto analytic_comparison          = sub_tag(solver & physics & residuals, "analytic_comparison");
    inline constexpr auto self_consistency             = sub_tag(solver & physics, "self_consistency");

    inline constexpr auto centrifugal                  = sub_tag(solver & physics, "centrifugal");
    inline constexpr auto advection                    = sub_tag(solver & physics, "advection");
    inline constexpr auto coriolis                     = sub_tag(solver & physics, "coriolis");
    inline constexpr auto gravity                      = sub_tag(solver & physics, "gravity");
    inline constexpr auto enthalpy                     = sub_tag(solver & physics, "enthalpy");
    inline constexpr auto barotrope                    = sub_tag(physics, "barotrope");
    inline constexpr auto mass_continuity              = sub_tag(solver & physics, "mass_continuity");
    inline constexpr auto pressure_gradient            = sub_tag(solver & physics, "pressure_gradient");
    inline constexpr auto viscosity                    = sub_tag(solver & physics, "viscosity");

    inline constexpr auto compactification             = sub_tag(mesh & mapping, "compactification");
    inline constexpr auto kelvin                       = sub_tag(compactification, "kelvin");
    inline constexpr auto mapping_evaluator            = mapping & make_tag("grid_function_evaluator");
    inline constexpr auto mapping_evaluator_unit       = mapping_evaluator & unit;

    inline constexpr auto prepared                     = sub_tag(solver & physics, "prepared");
    inline constexpr auto contexts                     = sub_tag(solver, "contexts");

    inline constexpr auto domain                       = sub_tag(mesh, "domain");

    // Canonical gravity-suite tags. These intentionally compose leaf tags
    // exactly once so Catch2 output remains useful and free of repeated
    // [solver]/[physics] entries inherited from older composite tags.
    inline constexpr auto gravity_unit                 = gravity & unit;
    inline constexpr auto gravity_integration          = gravity & integration;
    inline constexpr auto gravity_operator             = gravity & mfem_operators;
    inline constexpr auto gravity_prepared             = gravity & make_tag("prepared");
    inline constexpr auto gravity_context              = gravity & make_tag("context");
    inline constexpr auto gravity_kernel               = gravity & kernels;
    inline constexpr auto gravity_accuracy             = gravity & accuracy;
    inline constexpr auto gravity_operator_unit        = gravity_operator & unit;
    inline constexpr auto gravity_operator_integration = gravity_operator & integration;
    inline constexpr auto gravity_operator_convergence = gravity_operator & integration & make_tag("convergence");
    inline constexpr auto gravity_analytic             = gravity & integration & make_tag("analytic_comparison");
    inline constexpr auto gravity_consistency          = gravity & integration & make_tag("self_consistency");
    inline constexpr auto gravity_prepared_jacobian    = gravity_prepared & integration & make_tag("jacobian");
    inline constexpr auto gravity_prepared_unit        = gravity_prepared & unit;
    inline constexpr auto gravity_prepared_jacobian_accuracy = gravity_prepared_jacobian & accuracy;
    inline constexpr auto gravity_kernel_accuracy            = gravity_kernel & accuracy;
    inline constexpr auto gravity_kernel_integration         = gravity_kernel & integration;
    inline constexpr auto gravity_kernel_convergence         = gravity_kernel & integration & make_tag("convergence");
    inline constexpr auto gravity_analytic_accuracy          = gravity_analytic & accuracy;
    inline constexpr auto gravity_consistency_accuracy       = gravity_consistency & accuracy;
    inline constexpr auto gravity_integrator_unit            = gravity & integrator & unit;

    inline constexpr auto barotrope_prepared                 = barotrope & solver & make_tag("prepared");
    inline constexpr auto barotrope_eos_unit                 = barotrope & unit & make_tag("eos");
    inline constexpr auto barotrope_eos_jacobian             = barotrope_eos_unit & integration & make_tag("jacobian");
    inline constexpr auto polytropic_eos_characterization =
        barotrope & unit & make_tag("eos") & make_tag("characterization");
    inline constexpr auto polytropic_eos_relation_contract =
        barotrope & unit & make_tag("eos") & make_tag("relation_contract");
    inline constexpr auto polytropic_eos_compatibility = barotrope & unit & make_tag("eos") & make_tag("compatibility");
    inline constexpr auto equation_of_state            = physics & make_tag("eos");
    inline constexpr auto equation_of_state_type_system    = equation_of_state & unit & make_tag("type_system");
    inline constexpr auto equation_of_state_quantity_types = equation_of_state_type_system & make_tag("quantity_types");
    inline constexpr auto equation_of_state_relation_contract =
        equation_of_state_type_system & make_tag("relation_contract");
    inline constexpr auto equation_of_state_runtime_view = equation_of_state & unit & make_tag("runtime_view");
    inline constexpr auto equation_of_state_runtime_contract =
        equation_of_state_runtime_view & make_tag("relation_contract");
    inline constexpr auto equation_of_state_runtime_compatibility =
        equation_of_state_runtime_view & make_tag("compatibility");
    inline constexpr auto equation_of_state_consumer_contract =
        equation_of_state & unit & make_tag("consumer_contract");
    inline constexpr auto barotropic_closure_equation_of_state_contract =
        equation_of_state_consumer_contract & make_tag("barotropic_closure");
    inline constexpr auto pressure_force_equation_of_state_contract =
        equation_of_state_consumer_contract & make_tag("pressure_force");
    inline constexpr auto structure_seed_equation_of_state_contract =
        equation_of_state_consumer_contract & make_tag("structure_seed");
    inline constexpr auto stellar_model_type_contract = barotrope & model & unit & make_tag("type_contract");
    inline constexpr auto stellar_model_runtime_view  = barotrope & model & unit & make_tag("runtime_view");
    inline constexpr auto surface_prescription_type_contract =
        surface & physics & unit & make_tag("prescription") & make_tag("type_contract");
    inline constexpr auto surface_constraint_compilation =
        surface & physics & unit & make_tag("constraint_compilation");
    inline constexpr auto surface_constraint_jacobian = surface_constraint_compilation & jacobian;
    inline constexpr auto surface_constraint_lifetime = surface & model & unit & make_tag("constraint_lifetime");
    inline constexpr auto surface_boundary_dof_topology =
        surface & field_dof & integration & make_tag("boundary_topology");
    inline constexpr auto surface_row_replacement =
        surface & barotrope_prepared & integration & make_tag("row_replacement");
    inline constexpr auto translational_centering = geometry & solver & make_tag("translational_centering");
    inline constexpr auto translational_centering_topology =
        translational_centering & field_dof & integration & make_tag("point_topology");
    inline constexpr auto translational_centering_enforcement =
        translational_centering & barotrope_prepared & integration & make_tag("row_replacement");
    inline constexpr auto barotrope_pressure_quadrature = barotrope & mesh & geometry & solver & make_tag("pressure") &
                                                          make_tag("pressure_gradient") & make_tag("quadrature");
    inline constexpr auto barotrope_pressure_quadrature_unit     = barotrope_pressure_quadrature & unit;
    inline constexpr auto barotrope_pressure_quadrature_accuracy = barotrope_pressure_quadrature & accuracy;
    inline constexpr auto barotrope_prepared_jacobian   = barotrope_prepared & integration & make_tag("jacobian");
    inline constexpr auto barotrope_context             = barotrope & solver & make_tag("context");
    inline constexpr auto barotrope_context_integration = barotrope_context & integration;
    inline constexpr auto barotrope_prepared_analytic =
        barotrope_prepared & integration & make_tag("analytic_comparison");
    inline constexpr auto barotrope_prepared_jacobian_accuracy = barotrope_prepared_jacobian & accuracy;
    inline constexpr auto barotrope_prepared_jacobian_geometry = barotrope_prepared_jacobian & geometry;
    inline constexpr auto barotrope_prepared_jacobian_unit     = barotrope_prepared_jacobian & unit;

    // Canonical hydrostatic-suite tags. The leaf tags are composed directly
    // so inherited [physics]/[solver] tags appear only once.
    inline constexpr auto barotrope_hydrostatic                = barotrope & solver & make_tag("hydro");
    inline constexpr auto barotrope_hydrostatic_context        = barotrope_hydrostatic & make_tag("context");
    inline constexpr auto barotrope_hydrostatic_prepared       = barotrope_hydrostatic & make_tag("prepared");
    inline constexpr auto barotrope_hydrostatic_prepared_residual =
        barotrope_hydrostatic_prepared & integration & make_tag("residual");
    inline constexpr auto barotrope_hydrostatic_prepared_jacobian =
        barotrope_hydrostatic_prepared & integration & make_tag("jacobian");
    inline constexpr auto barotrope_hydrostatic_prepared_analytic =
        barotrope_hydrostatic_prepared & integration & make_tag("analytic_comparison");

    inline constexpr auto barotrope_mass_normalization          = barotrope & solver & make_tag("mass_normalization");
    inline constexpr auto barotrope_mass_normalization_context  = barotrope_mass_normalization & make_tag("context");
    inline constexpr auto barotrope_mass_normalization_prepared = barotrope_mass_normalization & make_tag("prepared");
    inline constexpr auto barotrope_mass_normalization_jacobian =
        barotrope_mass_normalization_prepared & integration & make_tag("jacobian");
    inline constexpr auto barotrope_mass_normalization_analytic =
        barotrope_mass_normalization_prepared & integration & make_tag("analytic_comparison");

    inline constexpr auto rotation_prepared          = centrifugal & make_tag("prepared");
    inline constexpr auto rotation_context           = centrifugal & make_tag("context");
    inline constexpr auto rotation_analytic          = centrifugal & integration & make_tag("analytic_comparison");
    inline constexpr auto rotation_context_unit      = rotation_context & unit;
    inline constexpr auto rotation_prepared_unit     = rotation_prepared & unit;
    inline constexpr auto rotation_prepared_jacobian = rotation_prepared & integration & make_tag("jacobian");
    inline constexpr auto rotation_prepared_jacobian_accuracy = rotation_prepared_jacobian & accuracy;
    inline constexpr auto rotation_kernel_accuracy            = centrifugal & kernels & accuracy;
    inline constexpr auto rotation_integrator_unit            = centrifugal & integrator & unit;
    inline constexpr auto rotation_integrator_integration     = centrifugal & integrator & integration;
    inline constexpr auto rotation_integrator_convergence =
        rotation_integrator_integration & convergence & h_refinement;
    inline constexpr auto rotation_analytic_unit              = rotation_analytic & unit;
    inline constexpr auto rotation_analytic_accuracy          = rotation_analytic & accuracy;
    inline constexpr auto rotation_analytic_accuracy_geometry = rotation_analytic_accuracy & geometry;

} // namespace tags