# Source code for dolfinx.plot

    # Copyright (C) 2021-2022 Jørgen S. Dokken and Garth N. Wells
    #
    # This file is part of DOLFINx (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    """Support functions for plotting"""

    import functools

    import numpy as np

    from dolfinx import cpp as _cpp
    from dolfinx import fem, mesh

    # NOTE: This dictionary and the below function that uses it should be
    # revised when pyvista improves rendering of 'arbitrary' Lagrange
    # elements, i.e. for the VTK cell types that define a shape but allow
    # arbitrary degree geometry. See
    # https://github.com/pyvista/pyvista/issues/947.
    #
    # Cell types can be found at
    # https://vtk.org/doc/nightly/html/vtkCellType_8h_source.html
    _first_order_vtk = {
        mesh.CellType.interval: 3,
        mesh.CellType.triangle: 5,
        mesh.CellType.quadrilateral: 9,
        mesh.CellType.tetrahedron: 10,
        mesh.CellType.hexahedron: 12,
    }

    @functools.singledispatch
    def vtk_mesh(msh: mesh.Mesh, dim: int | None = None, entities=None):
        """Create vtk mesh topology data for mesh entities of a given
        dimension. The vertex indices in the returned topology array are the
        indices for the associated entry in the mesh geometry.

        Args:
            mesh: Mesh to extract data from.
            dim: Topological dimension of entities to extract.
            entities: Entities to extract. Extract all if ``None``.

        Returns:
            Topology, type for each cell, and geometry in VTK-ready format.

        """
        if dim is None:
            dim = msh.topology.dim

        cell_type = _cpp.mesh.cell_entity_type(msh.topology.cell_type, dim, 0)
        if cell_type == mesh.CellType.prism:
            raise RuntimeError("Plotting of prism meshes not supported")

        # Use all local cells if not supplied
        if entities is None:
            entities = np.arange(msh.topology.index_map(dim).size_local, dtype=np.int32)

        geometry_entities = _cpp.mesh.entities_to_geometry(msh._cpp_object, dim, entities, False)

        num_nodes_per_cell = geometry_entities.shape[1]
        map_vtk = np.argsort(_cpp.io.perm_vtk(cell_type, num_nodes_per_cell))
        vtk_topology = geometry_entities[:, map_vtk]

        # Create mesh topology
        topology = np.empty((len(entities), num_nodes_per_cell + 1), dtype=np.int32)
        topology[:, 0] = num_nodes_per_cell
        topology[:, 1:] = vtk_topology

        # Array holding the cell type (shape) for each cell
        vtk_type = _cpp.io.get_vtk_cell_type(cell_type, dim)
        cell_types = np.full(len(entities), vtk_type)

        return topology.reshape(-1), cell_types, msh.geometry.x

    @vtk_mesh.register
    def _(V: fem.FunctionSpace, entities=None):  # type: ignore
        """Creates a VTK mesh topology (topology array and array of cell
        types) that is based on the degree-of-freedom coordinates.

        This function supports visualisation when the degree of the finite
        element space is different from the geometric degree of the mesh.

        Note:
            This function supports Lagrange elements (continuous and
            discontinuous) only.

        Args:
            V: Mesh to extract data from.
            entities: Entities to extract. Extract all if ``None``.

        Returns:
            Topology, type for each cell, and geometry in VTK-ready format.

        """
        if V.ufl_element().family_name not in [
            "Discontinuous Lagrange",
            "Lagrange",
            "DQ",
            "Q",
            "DP",
            "P",
        ]:
            raise RuntimeError(
                "Can only create meshes from continuous or discontinuous Lagrange"
                "spaces, not {V.ufl_element().family_name}."
            )

        degree = V.ufl_element().degree
        if degree == 0:
            raise RuntimeError("Cannot create topology from cellwise constants.")

        # Use all local cells if not supplied
        msh = V.mesh
        tdim = msh.topology.dim
        if entities is None:
            entities = range(msh.topology.index_map(tdim).size_local)

        dofmap = V.dofmap
        num_dofs_per_cell = V.dofmap.dof_layout.num_dofs
        cell_type = msh.topology.cell_type
        perm = np.argsort(_cpp.io.perm_vtk(cell_type, num_dofs_per_cell))

        vtk_type = (
            _first_order_vtk[cell_type] if degree == 1 else _cpp.io.get_vtk_cell_type(cell_type, tdim)
        )
        cell_types = np.full(len(entities), vtk_type)

        topology = np.zeros((len(entities), num_dofs_per_cell + 1), dtype=np.int32)
        topology[:, 0] = num_dofs_per_cell
        dofmap_ = dofmap.list
        topology[:, 1:] = dofmap_[entities][:, perm]
        return topology.reshape(1, -1)[0], cell_types, V.tabulate_dof_coordinates()