module;
#include <functional>
#include <string_view>

#include <mfem.hpp>

#include "xad_promote_polyfill.h"
#include <XAD/XAD.hpp>

export module mean_field:utils.misc;
import :utils.domain;

export namespace mean_field::utils {
    constexpr double APPROX_MAX_ACCEPTABLE_POTENTIAL_ERROR_SI_BURNING = 1e-4;

    bool is_vacuum(
        const mfem::ElementTransformation &Tr,
        mfem::Array<mfem::Vector *> elvec
    ) {
        using Schema = domain::CoreEnvelopeVacuumDomainSchema;

        if (Schema::template attribute_belongs_to<domain::Vacuum>(Tr.Attribute)) {
            const int size_elvec = elvec.Size();
            for (int i = 0; i < size_elvec; i++) {
                if (elvec[i]) {
                    *elvec[i] = 0.0;
                }
            }
            return true;
        }
        return false;
    }

    bool is_vacuum(
        const mfem::ElementTransformation &Tr,
        const mfem::Array2D<mfem::DenseMatrix *> &elmats
    ) {
        using Schema = domain::CoreEnvelopeVacuumDomainSchema;

        if (Schema::template attribute_belongs_to<domain::Vacuum>(Tr.Attribute)) {
            const int cols = elmats.NumCols();
            const int rows = elmats.NumRows();
            for (int rowID = 0; rowID < rows; rowID++) {
                for (int colID = 0; colID < cols; colID++) {
                    if (elmats(rowID, colID)) {
                        *elmats(rowID, colID) = 0.0;
                    }
                }
            }
            return true;
        }
        return false;
    }

    constexpr std::string_view ANSI_GREEN    = "\033[32m";
    constexpr std::string_view ANSI_RED      = "\033[31m";
    constexpr std::string_view ANSI_YELLOW   = "\033[33m";
    constexpr std::string_view ANSI_BLUE     = "\033[34m";
    constexpr std::string_view ANSI_MAGENTA  = "\033[35m";
    constexpr std::string_view ANSI_CYAN     = "\033[36m";
    constexpr std::string_view ANSI_RESET    = "\033[0m";
    constexpr std::string_view ANSI_BCYAN    = "\033[1;36m";

    constexpr double G                       = 1.0;
    constexpr double MASS                    = 1.0;
    constexpr double RADIUS                  = 1.0;

    [[maybe_unused]] constexpr char HOST[10] = "localhost";
    [[maybe_unused]] constexpr int PORT      = 19916;

    template <typename T>
    concept is_xad = std::is_same_v<T, xad::AReal<long double>> || std::is_same_v<T, xad::AReal<double>> ||
                     std::is_same_v<T, xad::AReal<float>>;

    template <typename T>
    concept is_real                  = std::is_floating_point_v<T> || is_xad<T>;

    template <is_real T> using EOS_P = std::function<T(const T &rho, const T &temp)>;

    enum class DOMAINS : uint8_t {
        CORE     = 1 << 0,
        ENVELOPE = 1 << 1,
        VACUUM   = 1 << 2,
        STELLAR  = CORE | ENVELOPE,
        ALL      = CORE | ENVELOPE | VACUUM
    };

    DOMAINS operator|(
        DOMAINS lhs,
        DOMAINS rhs
    );

    DOMAINS operator&(
        DOMAINS lhs,
        DOMAINS rhs
    );

    int get_mesh_order(const mfem::Mesh &mesh);

} // namespace mean_field::utils