#include "serif/discretization/domain/mesh/topology.hpp"

#include <mfem.hpp>

namespace serif::discretization::domain::mesh {
    const std::vector<int> MeshTopology::s_noElements{};

    namespace {
        [[nodiscard]] std::size_t as_index(const int id) noexcept {
            return static_cast<std::size_t>(id);
        }

        [[nodiscard]] bool in_range(const int id, const std::size_t size) noexcept {
            return id >= 0 && as_index(id) < size;
        }
    }

    MeshTopology::MeshTopology(const mfem::Mesh &mesh)
        : m_elementDomainIDs(as_index(mesh.GetNE()), 0),
          m_faceElements(as_index(mesh.GetNumFaces())),
          m_elementNeighbors(as_index(mesh.GetNE())),
          m_boundaryElementBoundaryIDs(as_index(mesh.GetNBE()), 0),
          m_boundaryElementFaceIDs(as_index(mesh.GetNBE()), -1),
          m_boundaryElementsByFace(as_index(mesh.GetNumFaces())) {

        // Which domain each volume element claims to belong to.
        for (int elementID = 0; elementID < mesh.GetNE(); ++elementID) {
            m_elementDomainIDs[as_index(elementID)] = mesh.GetAttribute(elementID);
        }

        // Face connectivity.
        for (int faceID = 0; faceID < mesh.GetNumFaces(); ++faceID) {
            int firstElementID  = -1;
            int secondElementID = -1;

            mesh.GetFaceElements(faceID, &firstElementID, &secondElementID);

            m_faceElements[as_index(faceID)] = FaceElements{
                .firstElementID  = firstElementID,
                .secondElementID = secondElementID
            };

            if (firstElementID < 0 || secondElementID < 0) {
                continue;
            }

            m_elementNeighbors[as_index(firstElementID)].push_back(secondElementID);
            m_elementNeighbors[as_index(secondElementID)].push_back(firstElementID);
        }

        // Boundary elements
        for (int boundaryElementID = 0; boundaryElementID < mesh.GetNBE(); ++boundaryElementID) {
            m_boundaryElementBoundaryIDs[as_index(boundaryElementID)] = mesh.GetBdrAttribute(boundaryElementID);

            const int faceID = mesh.GetBdrElementFaceIndex(boundaryElementID);

            m_boundaryElementFaceIDs[as_index(boundaryElementID)] = faceID;

            if (in_range(faceID, m_boundaryElementsByFace.size())) {
                m_boundaryElementsByFace[as_index(faceID)].push_back(boundaryElementID);
            }
        }
    }

    int MeshTopology::element_count() const noexcept {
        return static_cast<int>(m_elementDomainIDs.size());
    }

    int MeshTopology::face_count() const noexcept {
        return static_cast<int>(m_faceElements.size());
    }

    int MeshTopology::boundary_element_count() const noexcept {
        return static_cast<int>(m_boundaryElementBoundaryIDs.size());
    }

    int MeshTopology::element_domain_id(const int elementID) const {
        if (!in_range(elementID, m_elementDomainIDs.size())) {
            return 0;
        }
        return m_elementDomainIDs[as_index(elementID)];
    }

    int MeshTopology::boundary_element_boundary_id(const int boundaryElementID) const {
        if (!in_range(boundaryElementID, m_boundaryElementBoundaryIDs.size())) {
            return 0;
        }
        return m_boundaryElementBoundaryIDs[as_index(boundaryElementID)];
    }

    FaceElements MeshTopology::face_elements(const int faceID) const {
        if (!in_range(faceID, m_faceElements.size())) {
            return FaceElements{};
        }
        return m_faceElements[as_index(faceID)];
    }

    const std::vector<int> &MeshTopology::element_neighbors(const int elementID) const {
        if (!in_range(elementID, m_elementNeighbors.size())) {
            return s_noElements;
        }
        return m_elementNeighbors[as_index(elementID)];
    }

    int MeshTopology::boundary_element_face(const int boundaryElementID) const {
        if (!in_range(boundaryElementID, m_boundaryElementFaceIDs.size())) {
            return -1;
        }
        return m_boundaryElementFaceIDs[as_index(boundaryElementID)];
    }

    const std::vector<int> &MeshTopology::boundary_elements_on_face(const int faceID) const {
        if (!in_range(faceID, m_boundaryElementsByFace.size())) {
            return s_noElements;
        }
        return m_boundaryElementsByFace[as_index(faceID)];
    }
}