# Source code for dolfinx.io.gmsh

    # Copyright (C) 2022-2025 Jørgen S. Dokken, Henrik N. T. Finsberg and
    # Paul T. Kühner
    #
    # This file is part of DOLFINx (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    """Tools to extract data from Gmsh models."""

    import typing
    from collections.abc import Callable
    from pathlib import Path

    from mpi4py import MPI as _MPI

    import numpy as np
    import numpy.typing as npt

    import basix
    import basix.ufl
    import ufl
    from dolfinx import cpp as _cpp
    from dolfinx import default_real_type
    from dolfinx.cpp.graph import AdjacencyList_int32 as _AdjacencyList_int32
    from dolfinx.graph import AdjacencyList, adjacencylist
    from dolfinx.io.utils import distribute_entity_data
    from dolfinx.mesh import CellType, Mesh, MeshTags, create_mesh, meshtags_from_entities

    __all__ = [
        "MeshData",
        "PhysicalGroup",
        "TopologyDict",
        "cell_perm_array",
        "extract_geometry",
        "extract_topology_and_markers",
        "model_to_mesh",
        "read_from_msh",
        "ufl_mesh",
    ]

    class TopologyDict(typing.TypedDict):
        """TopologyDict is a TypedDict for storing the topology of the marked
        cell.

        Args:
            topology: 2D array containing the topology of the marked cell.
            cell_data: List with the corresponding markers.

        Note:
            The TypedDict is only used for type hinting, and does not
            enforce the structure of the dictionary, but rather provides
            a hint to the user and the type checker.
        """

        topology: npt.NDArray[
            np.integer
        ]  # cell->node connectivity, shape (num_cells, num_nodes_per_cell)
        cell_data: npt.NDArray[np.integer]  # cell tags, shape (num_cells,)
        entity_tags: npt.NDArray[np.integer]  # GMSH cell indices, shape (num_cells,)

    # Map from Gmsh cell type identifier (integer) to DOLFINx cell type and
    # degree https://gmsh.info//doc/texinfo/gmsh.html#MSH-file-format
    _gmsh_to_cells = {
        1: ("interval", 1),
        2: ("triangle", 1),
        3: ("quadrilateral", 1),
        4: ("tetrahedron", 1),
        5: ("hexahedron", 1),
        8: ("interval", 2),
        9: ("triangle", 2),
        10: ("quadrilateral", 2),
        11: ("tetrahedron", 2),
        12: ("hexahedron", 2),
        15: ("point", 0),
        21: ("triangle", 3),
        26: ("interval", 3),
        29: ("tetrahedron", 3),
        36: ("quadrilateral", 3),
        92: ("hexahedron", 3),
    }

    class PhysicalGroup(typing.NamedTuple):
        """Physical group info.

        Args:
            dim: dimension of the physical group
            tag: tag of the physical group
        """

        dim: int
        tag: int

    class MeshData(typing.NamedTuple):
        """Data for representing a mesh and associated tags.

        Args:
            mesh: Mesh.
            cell_tags: MeshTags for cells.
            facet_tags: MeshTags for facets (codim 1).
            ridge_tags: MeshTags for ridges (codim 2).
            peak_tags: MeshTags for peaks (codim 3).
            physical_groups: Physical groups in the mesh, where the key
                is the physical name and the value is a tuple with the
                dimension and tag.
        """

        mesh: Mesh
        cell_tags: MeshTags | None
        facet_tags: MeshTags | None
        ridge_tags: MeshTags | None
        peak_tags: MeshTags | None
        physical_groups: dict[str, PhysicalGroup]

    def ufl_mesh(gmsh_cell: int, gdim: int, dtype: npt.DTypeLike) -> ufl.Mesh:
        """Create a UFL mesh from a Gmsh cell identifier and geometric
        dimension.

        See https://gmsh.info//doc/texinfo/gmsh.html#MSH-file-format.

        Args:
            gmsh_cell: Gmsh cell identifier.
            gdim: Geometric dimension of the mesh.

        Returns:
            UFL Mesh using Lagrange elements (equispaced) of the
            corresponding DOLFINx cell.

        """
        try:
            shape, degree = _gmsh_to_cells[gmsh_cell]
        except KeyError as e:
            print(f"Unknown cell type {gmsh_cell}.")
            raise e
        cell = ufl.Cell(shape)
        element = basix.ufl.element(
            basix.ElementFamily.P,
            cell.cellname(),
            degree,
            basix.LagrangeVariant.equispaced,
            shape=(gdim,),
            dtype=dtype,  # type: ignore[arg-type]
        )
        return ufl.Mesh(element)

    def cell_perm_array(cell_type: CellType, num_nodes: int) -> list[int]:
        """The permutation array for permuting Gmsh ordering to DOLFINx
        ordering.

        Args:
            cell_type: DOLFINx cell type.
            num_nodes: Number of nodes in the cell.

        Returns:
            An array ``p`` such that ``a_dolfinx[i] = a_gmsh[p[i]]``.

        """
        return _cpp.io.perm_gmsh(cell_type, num_nodes)

    def extract_topology_and_markers(
        model, name: str | None = None
    ) -> tuple[dict[int, TopologyDict], dict[str, PhysicalGroup]]:
        """Extract all entities tagged with a physical marker in the gmsh
        model.

        Returns a nested dictionary where the first key is the gmsh MSH
        element type integer. Each element type present in the model
        contains the cell topology of the elements and corresponding
        markers.

        Args:
            model: Gmsh model.
            name: Name of the gmsh model. If not set the current
                model will be used.

        Returns:
            A tuple ``(topologies, physical_groups)``, where ``topologies`` is
            a nested dictionary where each key corresponds to a gmsh cell type.
            Each cell type found in the mesh has a 2D array containing the
            topology of the marked cell and a list with the corresponding
            markers. ``physical_groups`` is a dictionary where the key is the
            physical name and the value is a tuple with the dimension and tag.

        """
        if name is not None:
            model.setCurrent(name)

        # Get the physical groups from gmsh in the form [(dim1, tag1),
        # (dim1, tag2), (dim2, tag3),...]
        phys_grps = model.getPhysicalGroups()
        topologies: dict[int, TopologyDict] = {}
        # Create a dictionary with the physical groups where the key is the
        # physical name and the value is a tuple with the dimension and tag
        physical_groups: dict[str, PhysicalGroup] = {}
        for dim, tag in phys_grps:
            # Get the entities of dimension `dim`, dim=0 -> Points, dim=1 -
            # >Lines, dim=2 -> Triangles/Quadrilaterals, etc.
            entities = model.getEntitiesForPhysicalGroup(dim, tag)
            for entity in entities:
                # Get cell type, list of cells with given tag and topology
                # of tagged cells NOTE: Assumes that each entity only have
                # one cell-type,
                # i.e. facets of prisms and pyramid meshes are not supported
                (entity_types, entity_tags, entity_topologies) = model.mesh.getElements(dim, tag=entity)

                if len(entity_types) > 1:
                    raise RuntimeError(
                        f"Unsupported mesh with multiple cell types {entity_types} for entity {entity}"
                    )
                elif len(entity_types) == 0:
                    continue

                # Determine number of local nodes per element to create the
                # topology of the elements
                properties = model.mesh.getElementProperties(entity_types[0])
                name, dim, _, num_nodes, _, _ = properties

                # Array of shape (num_elements,num_nodes_per_element)
                # containing the topology of the elements on this entity.
                # NOTE: Gmsh indexing starts with one, we therefore subtract
                # 1 from each node to use zero-based numbering
                topology = entity_topologies[0].reshape(-1, num_nodes) - 1

                # Create marker array of length of number of tagged cells
                marker = np.full_like(entity_tags[0], tag)

                # Keep track of entity tags to check tag consistency
                entity_tags = entity_tags[0]

                # Group element topology and markers of the same entity type
                entity_type = entity_types[0]
                if entity_type in topologies.keys():
                    topologies[entity_type]["topology"] = np.concatenate(
                        (topologies[entity_type]["topology"], topology), axis=0
                    )
                    topologies[entity_type]["cell_data"] = np.hstack(
                        [topologies[entity_type]["cell_data"], marker]
                    )
                    topologies[entity_type]["entity_tags"] = np.hstack(
                        [topologies[entity_type]["entity_tags"], entity_tags]
                    )
                else:
                    topologies[entity_type] = {
                        "topology": topology,
                        "cell_data": marker,
                        "entity_tags": entity_tags,
                    }
            physical_groups[model.getPhysicalName(dim, tag)] = PhysicalGroup(dim, tag)

        return topologies, physical_groups

    def extract_geometry(model, name: str | None = None) -> npt.NDArray[np.float64]:
        """Extract the mesh geometry from a Gmsh model.

        Returns an array of shape ``(num_nodes, 3)``, where the i-th row
        corresponds to the i-th node in the mesh.

        Args:
            model: Gmsh model
            name: Name of the Gmsh model. If not set the current
                model will be used.

        Returns:
            The mesh geometry as an array of shape ``(num_nodes, 3)``.

        """
        if name is not None:
            model.setCurrent(name)

        # Get the unique tag and coordinates for nodes in mesh
        indices, points, _ = model.mesh.getNodes()
        points = points.reshape(-1, 3)

        # Gmsh indices starts at 1. We therefore subtract one to use
        # zero-based numbering
        indices -= 1

        # In some cases, Gmsh does not return the points in the same
        # order as their unique node index. We therefore sort nodes in
        # geometry according to the unique index
        perm_sort = np.argsort(indices)
        assert np.all(indices[perm_sort] == np.arange(len(indices)))
        return points[perm_sort]

    def model_to_mesh(
        model,
        comm: _MPI.Comm,
        rank: int,
        gdim: int = 3,
        partitioner: Callable[[_MPI.Comm, int, int, _AdjacencyList_int32], _AdjacencyList_int32]
        | None = None,
        dtype=default_real_type,
    ) -> MeshData:
        """Create a Mesh from a Gmsh model.

        Creates a :class:`dolfinx.mesh.Mesh` from the physical entities of
        the highest topological dimension in the Gmsh model. In parallel,
        the gmsh model is processed on one MPI rank, and the
        :class:`dolfinx.mesh.Mesh` is distributed across ranks.

        Args:
            model: Gmsh model.
            comm: MPI communicator to use for mesh creation.
            rank: MPI rank that the Gmsh model is initialized on.
            gdim: Geometrical dimension of the mesh.
            partitioner: Function that computes the parallel
                distribution of cells across MPI ranks.

        Returns:
            MeshData with mesh and tags of corresponding entities by
            codimension. Codimension 0 is the cell tags, codimension 1 is the
            facet tags, codimension 2 is the ridge tags and codimension 3 is
            the peak tags as well as a lookup table from the physical groups by
            name to integer.

        Note:
            For performance, this function should only be called once for
            large problems. For reuse, it is recommended to save the mesh
            and corresponding tags using :class:`dolfinx.io.XDMFFile` after
            creation for efficient access.
        """
        valid_mesh = None
        if comm.rank == rank:
            assert model is not None, "Gmsh model is None on rank responsible for mesh creation."
            # Get mesh geometry and mesh topology for each element
            x = extract_geometry(model)
            topologies, physical_groups = extract_topology_and_markers(model)
            valid_mesh = len(physical_groups) > 0

            # Extract Gmsh entity (cell) id, topological dimension and number
            # of nodes which is used to create an appropriate coordinate
            # element, and seperate higher topological entities from lower
            # topological entities (e.g. facets, ridges and peaks).
            num_unique_entities = len(topologies.keys())
            element_ids = np.zeros(num_unique_entities, dtype=np.int32)
            entity_tdim = np.zeros(num_unique_entities, dtype=np.int32)
            num_nodes_per_element = np.zeros(num_unique_entities, dtype=np.int32)
            for i, element in enumerate(topologies.keys()):
                _, dim, _, num_nodes, _, _ = model.mesh.getElementProperties(element)
                element_ids[i] = element
                entity_tdim[i] = dim
                num_nodes_per_element[i] = num_nodes

            # Broadcast information to all other ranks
            entity_tdim, element_ids, num_nodes_per_element = comm.bcast(
                (entity_tdim, element_ids, num_nodes_per_element), root=rank
            )
        else:
            entity_tdim, element_ids, num_nodes_per_element = comm.bcast((None, None, None), root=rank)

        valid_mesh = comm.bcast(valid_mesh, root=rank)
        if not valid_mesh:
            raise RuntimeError("No 'physical groups' in gmsh mesh. Cannot continue.")

        # Sort elements by descending dimension
        assert len(np.unique(entity_tdim)) == len(entity_tdim)
        perm_sort = np.argsort(entity_tdim)[::-1]

        # Extract position of the highest topological entity and its
        # topological dimension
        cell_position = perm_sort[0]
        tdim = int(entity_tdim[cell_position])

        # Check that all cells are tagged once
        error_msg = ""
        if comm.rank == rank:
            _elementTypes, _elementTags, _nodeTags = model.mesh.getElements(dim=tdim, tag=-1)
            _elementType_dim = _elementTypes[0]
            if _elementType_dim not in topologies.keys():
                error_msg = "All cells are expected to be tagged once; none found"

            num_cells_tagged = len(topologies[_elementType_dim]["entity_tags"])
            if (num_cells := len(_elementTags[0])) != num_cells_tagged:
                error_msg = (
                    "All cells are expected to be tagged once;"
                    + f"found: {num_cells_tagged}, expected: {num_cells}"
                )
            num_cells_tagged_once = len(np.unique(topologies[_elementType_dim]["entity_tags"]))
            if num_cells_tagged != num_cells_tagged_once:
                error_msg = "All cells are expected to be tagged once, found duplicates"

        error_msg = comm.bcast(error_msg, root=rank)
        if error_msg != "":
            raise RuntimeError(error_msg)

        # Extract entity -> node connectivity for all cells and sub-entities
        # marked in the GMSH model
        meshtags: dict[int, tuple[npt.NDArray[np.int64], npt.NDArray[np.int32]]] = {}
        for position in perm_sort:
            codim = tdim - entity_tdim[position]
            if comm.rank == rank:
                gmsh_entity_id = element_ids[position]
                marked_entities = np.asarray(topologies[gmsh_entity_id]["topology"], dtype=np.int64)
                entity_values = np.asarray(topologies[gmsh_entity_id]["cell_data"], dtype=np.int32)
                meshtags[codim] = (marked_entities, entity_values)
            else:
                # Any other process than input rank does not have any entities
                marked_entities = np.empty((0, num_nodes_per_element[position]), dtype=np.int32)
                entity_values = np.empty((0,), dtype=np.int32)
                meshtags[codim] = (marked_entities, entity_values)

        # Create a UFL Mesh object for the GMSH element with the highest
        # topoligcal dimension
        ufl_domain = ufl_mesh(element_ids[cell_position], gdim, dtype=dtype)

        # Get cell->node connectivity and  permute to FEniCS ordering
        num_nodes = num_nodes_per_element[cell_position]
        gmsh_cell_perm = cell_perm_array(_cpp.mesh.to_type(str(ufl_domain.ufl_cell())), num_nodes)
        cell_connectivity = meshtags[0][0][:, gmsh_cell_perm].copy()

        # Create a distributed mesh, where mesh nodes are only destributed from
        # the input rank
        if comm.rank != rank:
            x = np.empty([0, gdim], dtype=dtype)  # No nodes on other than root rank
        mesh = create_mesh(
            comm, cell_connectivity, ufl_domain, x[:, :gdim].astype(dtype, copy=False), partitioner
        )
        assert tdim == mesh.topology.dim, (
            f"{mesh.topology.dim=} does not match Gmsh model dimension {tdim}"
        )

        # Create MeshTags for all sub entities
        topology = mesh.topology
        codim_to_name = {0: "cell", 1: "facet", 2: "ridge", 3: "peak"}
        dolfinx_meshtags: dict[str, MeshTags | None] = {}
        for codim in [0, 1, 2, 3]:
            key = f"{codim_to_name[codim]}_tags"
            if (
                codim == 1 and topology.cell_type == CellType.prism
            ) or topology.cell_type == CellType.pyramid:
                raise RuntimeError(f"Unsupported facet tag for type {topology.cell_type}")

            meshtag_data = meshtags.get(codim, None)
            if meshtag_data is None:
                dolfinx_meshtags[key] = None
                continue

            # Distribute entity data [[e0_v0, e0_v1, ...], [e1_v0, e1_v1, ...],
            # ...] which is made in global input indices to local indices on
            # the owning process
            (marked_entities, entity_values) = meshtag_data
            local_entities, local_values = distribute_entity_data(
                mesh, tdim - codim, marked_entities, entity_values
            )
            # Create MeshTags object from the local entities
            mesh.topology.create_connectivity(tdim - codim, tdim)
            adj = adjacencylist(local_entities)
            et = meshtags_from_entities(
                mesh, tdim - codim, adj, local_values.astype(np.int32, copy=False)
            )
            et.name = key
            dolfinx_meshtags[key] = et

        # Broadcast physical groups (string to integer mapping) to all ranks
        if comm.rank == rank:
            physical_groups = comm.bcast(physical_groups, root=rank)
        else:
            physical_groups = comm.bcast(None, root=rank)

        return MeshData(mesh, physical_groups=physical_groups, **dolfinx_meshtags)

    def read_from_msh(
        filename: str | Path,
        comm: _MPI.Comm,
        rank: int = 0,
        gdim: int = 3,
        partitioner: Callable[[_MPI.Comm, int, int, AdjacencyList], _AdjacencyList_int32] | None = None,
    ) -> MeshData:
        """Read a Gmsh .msh file and return a :class:`dolfinx.mesh.Mesh` and
        cell facet markers.

        Note:
            This function requires the Gmsh Python module.

        Args:
            filename: Name of ``.msh`` file.
            comm: MPI communicator to create the mesh on.
            rank: Rank of ``comm`` responsible for reading the ``.msh``
                file.
            gdim: Geometric dimension of the mesh

        Returns:
            Meshdata with mesh, cell tags, facet tags, edge tags,
            vertex tags and physical groups.

        """
        try:
            import gmsh
        except ModuleNotFoundError:
            # Python 3.11+ adds the add_note method to exceptions
            # e.add_note("Gmsh must be installed to import dolfinx.io.gmsh")
            raise ModuleNotFoundError(
                "No module named 'gmsh': dolfinx.io.gmsh.read_from_msh requires Gmsh.", name="gmsh"
            )

        if comm.rank == rank:
            gmsh.initialize()
            gmsh.model.add("Mesh from file")
            gmsh.merge(str(filename))
            msh = model_to_mesh(gmsh.model, comm, rank, gdim=gdim, partitioner=partitioner)
            gmsh.finalize()
            return msh
        else:
            return model_to_mesh(gmsh.model, comm, rank, gdim=gdim, partitioner=partitioner)