module;
#include <mfem.hpp>
#include <expected>

module mean_field;
import :boundary.contexts;

namespace mean_field::utils {
    DOMAINS operator|(
        DOMAINS lhs,
        DOMAINS rhs
    ) {
        return static_cast<DOMAINS>(static_cast<uint8_t>(lhs) | static_cast<uint8_t>(rhs));
    }

    DOMAINS operator&(
        DOMAINS lhs,
        DOMAINS rhs
    ) {
        return static_cast<DOMAINS>(static_cast<uint8_t>(lhs) & static_cast<uint8_t>(rhs));
    }

    void populate_element_mask(
        const mfem::Mesh* mesh,
        const DOMAINS domain,
        mfem::Array<int> &mask
    ) {
        const int max_attr = mesh->attributes.Max();
        mask.SetSize(max_attr);
        mask = 0;

        if ((domain & DOMAINS::CORE) == DOMAINS::CORE && max_attr >= 1) {
            mask[0] = 1;
        }

        if ((domain & DOMAINS::ENVELOPE) == DOMAINS::ENVELOPE && max_attr >= 2) {
            mask[1] = 1;
        }

        if ((domain & DOMAINS::VACUUM) == DOMAINS::VACUUM && max_attr >= 3) {
            mask[2] = 1;
        }
    }

    void populate_domain_tdofs(
        const mfem::ParFiniteElementSpace *fes,
        const mfem::Array<int> &element_mask,
        mfem::Array<int> &ess_tdof
    ) {
        mfem::Array<int> vdof_marker(fes->GetVSize());
        vdof_marker = 0;

        for (int i = 0; i < fes->GetMesh()->GetNE(); i++) {
            const int attr = fes->GetMesh()->GetAttribute(i);

            if (element_mask[attr - 1]) {
                mfem::Array<int> dofs;
                fes->GetElementVDofs(i, dofs);

                for (int j = 0; j < dofs.Size(); j++) {
                    int index = dofs[j];
                    if (index < 0) index = -1 - index;
                    vdof_marker[index] = 1;
                }
            }
        }

        fes->MarkerToList(vdof_marker, ess_tdof);
    }


    std::expected<boundary::Bounds, boundary::BoundsError> discover_bounds(
        const mfem::Mesh *mesh,
        const int vacuum_attr
    ) {
        double local_min_r = std::numeric_limits<double>::max();
        double local_max_r = -std::numeric_limits<double>::max();
        bool found_vacuum = false;

        for (int i = 0; i < mesh->GetNE(); ++i) {
            if (mesh->GetAttribute(i) == vacuum_attr) {
                found_vacuum = true;
                mfem::Array<int> vertices;
                mesh->GetElementVertices(i, vertices);
                for (const int v: vertices) {
                    const double *coords = mesh->GetVertex(v);
                    double r = std::sqrt(coords[0] * coords[0] + coords[1] * coords[1] + coords[2] * coords[2]);
                    local_min_r = std::min(local_min_r, r);
                    local_max_r = std::max(local_max_r, r);
                }
            }
        }

        double global_min_r, global_max_r;
        int global_found_vacuum;
        int l_found = found_vacuum ? 1 : 0;

        MPI_Comm comm = MPI_COMM_WORLD;
        if (const auto *pmesh = dynamic_cast<const mfem::ParMesh *>(mesh)) {
            comm = pmesh->GetComm();
        }

        MPI_Allreduce(&local_min_r, &global_min_r, 1, MPI_DOUBLE, MPI_MIN, comm);
        MPI_Allreduce(&local_max_r, &global_max_r, 1, MPI_DOUBLE, MPI_MAX, comm);
        MPI_Allreduce(&l_found, &global_found_vacuum, 1, MPI_INT, MPI_MAX, comm);

        if (global_found_vacuum) {
            return boundary::Bounds(global_min_r, global_max_r);
        }
        return std::unexpected(boundary::BoundsError::CANNOT_FIND_VACUUM);
    }

    int get_mesh_order(
        const mfem::Mesh &mesh
    ) {
        if (mesh.GetNodes() != nullptr) {
            return mesh.GetNodes()->FESpace()->GetMaxElementOrder();
        }
        return 1;
    }



}