module;
#include <algorithm>
#include <array>
#include <catch2/internal/catch_stringref.hpp>
#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 {
    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;

        const int vacuum_attribute = f.domainMapperStateless->GetVacuumElementAttribute();

        for (int i = 0; i < f.mesh->attributes.Size(); ++i) {
            const int attribute   = f.mesh->attributes[i];
            const bool is_stellar = attribute != vacuum_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 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 legacy_comparison   = make_tag("legacy_comparison");
    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 prepared            = sub_tag(solver & physics, "prepared");
    inline constexpr auto contexts            = sub_tag(solver, "contexts");

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

} // namespace tags