# Source code for dolfinx.io.utils

    # Copyright (C) 2017-2022 Chris N. Richardson, Garth N. Wells,
    # Michal Habera and Jørgen S. Dokken
    #
    # This file is part of DOLFINx (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    """IO module for input data and post-processing file output."""

    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.cpp.io import perm_gmsh as cell_perm_gmsh
    from dolfinx.cpp.io import perm_vtk as cell_perm_vtk
    from dolfinx.fem import Function
    from dolfinx.mesh import CellType, Geometry, GhostMode, Mesh, MeshTags

    __all__ = ["VTKFile", "XDMFFile", "cell_perm_gmsh", "cell_perm_vtk", "distribute_entity_data"]

    # VTXWriter requires ADIOS2
    if _cpp.common.has_adios2:
        from dolfinx.cpp.io import VTXMeshPolicy  # F401

        __all__ = [*__all__, "VTXWriter", "VTXMeshPolicy"]

        class VTXWriter:
            """Writer for VTX files, using ADIOS2 to create the files.

            VTX supports arbitrary order Lagrange finite elements for the
            geometry description and arbitrary order (discontinuous)
            Lagrange finite elements for Functions.

            The files can be viewed using Paraview.
            """

            _cpp_object: _cpp.io.VTXWriter_float32 | _cpp.io.VTXWriter_float64

            def __init__(
                self,
                comm: _MPI.Comm,
                filename: str | Path,
                output: Mesh | Function | list[Function] | tuple[Function],
                engine: str = "BPFile",
                mesh_policy: VTXMeshPolicy = VTXMeshPolicy.update,
            ):
                """Initialize a writer for outputting data in the VTX format.

                Args:
                    comm: The MPI communicator
                    filename: The output filename
                    output: The data to output. Either a mesh, a single
                        (discontinuous) Lagrange Function or list of
                        (discontinuous) Lagrange Functions.
                    engine: ADIOS2 engine to use for output. See
                        ADIOS2 documentation for options.
                    mesh_policy: Controls if the mesh is written to file at
                        the first time step only when a ``Function`` is
                        written to file, or is re-written (updated) at each
                        time step. Has an effect only for ``Function``
                        output.

                Note:
                    All Functions for output must share the same mesh and
                    have the same element type.
                """
                # Get geometry type
                if isinstance(output, Mesh):
                    dtype = output.geometry.x.dtype
                elif isinstance(output, Function):
                    dtype = output.function_space.mesh.geometry.x.dtype
                else:
                    dtype = output[0].function_space.mesh.geometry.x.dtype

                if np.issubdtype(dtype, np.float32):
                    _vtxwriter = _cpp.io.VTXWriter_float32
                elif np.issubdtype(dtype, np.float64):
                    _vtxwriter = _cpp.io.VTXWriter_float64
                else:
                    raise RuntimeError(f"VTXWriter does not support dtype={dtype}.")

                if isinstance(output, Mesh):
                    self._cpp_object = _vtxwriter(comm, filename, output._cpp_object, engine)  # type: ignore[union-attr]
                else:
                    cpp_objects = (
                        [output._cpp_object]
                        if isinstance(output, Function)
                        else [o._cpp_object for o in output]
                    )
                    self._cpp_object = _vtxwriter(comm, filename, cpp_objects, engine, mesh_policy)

            def __enter__(self):
                return self

            def __exit__(self, exception_type, exception_value, traceback):
                self.close()

            def write(self, t: float):
                self._cpp_object.write(t)

            def close(self):
                self._cpp_object.close()

    class VTKFile(_cpp.io.VTKFile):
        """Interface to VTK files.

        VTK supports arbitrary order Lagrange finite elements for the
        geometry description. XDMF is the preferred format for geometry
        order <= 2.
        """

        def __enter__(self):
            return self

        def __exit__(self, exception_type, exception_value, traceback):
            self.close()

        def write_mesh(self, mesh: Mesh, t: float = 0.0) -> None:
            """Write mesh to file for a given time (default 0.0)"""
            self.write(mesh._cpp_object, t)

        def write_function(self, u: list[Function] | Function, t: float = 0.0) -> None:
            """Write a single function or a list of functions to file for a
            given time (default 0.0)"""
            cpp_objects = [u._cpp_object] if isinstance(u, Function) else [_u._cpp_object for _u in u]
            super().write(cpp_objects, t)

    class XDMFFile(_cpp.io.XDMFFile):
        Encoding = _cpp.io.XDMFFile.Encoding

        def __enter__(self):
            return self

        def __exit__(self, exception_type, exception_value, traceback):
            self.close()

        def write_mesh(self, mesh: Mesh, xpath: str = "/Xdmf/Domain") -> None:
            """Write mesh to file"""
            super().write_mesh(mesh._cpp_object, xpath)

        def write_meshtags(
            self,
            tags: MeshTags,
            x: Geometry,
            geometry_xpath: str = "/Xdmf/Domain/Grid/Geometry",
            xpath: str = "/Xdmf/Domain",
        ) -> None:
            """Write mesh tags to file"""
            super().write_meshtags(tags._cpp_object, x._cpp_object, geometry_xpath, xpath)

        def write_function(
            self, u: Function, t: float = 0.0, mesh_xpath="/Xdmf/Domain/Grid[@GridType='Uniform'][1]"
        ):
            """Write function to file for a given time.

            Note:
                Function is interpolated onto the mesh nodes, as a Nth order
                Lagrange function, where N is the order of the coordinate
                map. If the Function is a cell-wise constant, it is saved as
                a cell-wise constant.

            Args:
                u: Function to write to file.
                t: Time associated with Function output.
                mesh_xpath: Path to mesh associated with the Function in the
                    XDMFFile.
            """
            super().write_function(getattr(u, "_cpp_object", u), t, mesh_xpath)

        def read_mesh(
            self,
            ghost_mode=GhostMode.shared_facet,
            name="mesh",
            xpath="/Xdmf/Domain",
            max_facet_to_cell_links: int = 2,
        ) -> Mesh:
            """Read mesh data from file.

            Note:
                Changing `max_facet_to_cell_links` from the default value
                should only be required when working on branching manifolds.
                Changing this value on non-branching meshes will only result in
                a slower mesh partitioning and creation.

            Args:
                ghost_mode: Ghost mode to use for the cells in mesh creation.
                name: Name of the grid node in the xml-scheme in the file
                xpath: XPath where Mesh Grid is stored in the file.
                max_facet_to_cell_links: Maximum number of cells that a facet
                    can be linked to.
            """
            cell_shape, cell_degree = super().read_cell_type(name, xpath)
            cells = super().read_topology_data(name, xpath)
            x = super().read_geometry_data(name, xpath)

            # Get coordinate element, special handling for second order
            # serendipity.
            num_nodes_per_cell = cells.shape[1]
            if (cell_shape == CellType.quadrilateral and num_nodes_per_cell == 8) or (
                cell_shape == CellType.hexahedron and num_nodes_per_cell == 20
            ):
                s_el = basix.ufl.element(
                    basix.ElementFamily.serendipity,
                    cell_shape.name,
                    2,
                )
                # Create a custom element that is serendipity but uses points
                # evaluations on edges
                geometry = basix.cell.geometry(s_el.basix_element.cell_type)
                topology = basix.cell.topology(s_el.basix_element.cell_type)
                e_x: list[list[npt.NDArray[np.floating]]] = [
                    [np.array([p]) for p in geometry],
                    [np.array([(geometry[edge[0]] + geometry[edge[1]]) / 2]) for edge in topology[1]],
                    [np.zeros((0, 3)) for _ in s_el.basix_element.x[2]],
                    [np.zeros((0, 3)) for _ in s_el.basix_element.x[3]],
                ]
                e_m: list[list[npt.NDArray[np.floating]]] = [
                    [np.ones((1, 1, 1, 1)) for _ in s_el.basix_element.M[0]],
                    [np.ones((1, 1, 1, 1)) for _ in s_el.basix_element.M[1]],
                    [np.zeros((0, 1, 0, 1)) for _ in s_el.basix_element.M[2]],
                    [np.zeros((0, 1, 0, 1)) for _ in s_el.basix_element.M[3]],
                ]
                el = basix.ufl.custom_element(
                    s_el.basix_element.cell_type,
                    s_el.reference_value_shape,
                    s_el.basix_element.wcoeffs,
                    e_x,
                    e_m,
                    0,
                    s_el.map_type,
                    s_el.basix_element.sobolev_space,
                    s_el.discontinuous,
                    s_el.embedded_subdegree,
                    s_el.embedded_superdegree,
                    s_el.polyset_type,
                    s_el.dtype,
                )
                cmap = _cpp.fem.CoordinateElement_float64(el.basix_element._e)
                basix_el = basix.ufl.blocked_element(el, shape=(x.shape[1],))
            else:
                basix_el = basix.ufl.element(
                    "Lagrange",
                    cell_shape.name,
                    cell_degree,
                    basix.LagrangeVariant.unset,
                    shape=(x.shape[1],),
                )
                cmap = _cpp.fem.CoordinateElement_float64(cell_shape, cell_degree)

            # Build the mesh
            msh = _cpp.mesh.create_mesh(
                self.comm,
                cells,
                cmap,
                x,
                _cpp.mesh.create_cell_partitioner(ghost_mode),
                max_facet_to_cell_links,
            )
            msh.name = name
            domain = ufl.Mesh(basix_el)
            return Mesh(msh, domain)

        def read_meshtags(
            self,
            mesh: Mesh,
            name: str,
            attribute_name: str | None = None,
            xpath: str = "/Xdmf/Domain",
        ) -> MeshTags:
            """Read MeshTags with a specific name as specified in the XMDF
            file.

            Args:
                mesh: Mesh that the input data is defined on.
                name: Name of the grid node in the xml-scheme of the
                    XDMF-file.
                attribute_name: The name of the attribute to read. If
                    ``attribute_name`` is empty, reads the first attribute in
                    the file. If ``attribute_name`` is not empty but no
                    attributes have the provided name, throws an error. If
                    multiple attributes have the provided name, reads the
                    first one found.
                xpath: XPath where MeshTags Grid is stored in file.

            Returns:
                A MeshTags object containing the requested data read from
                file.
            """
            mt = super().read_meshtags(mesh._cpp_object, name, attribute_name, xpath)
            return MeshTags(mt)

    def distribute_entity_data(
        mesh: Mesh, entity_dim: int, entities: npt.NDArray[np.int64], values: np.ndarray
    ) -> tuple[npt.NDArray[np.int64], np.ndarray]:
        """Given a set of mesh entities and values, distribute them to the
        process that owns the entity.

        The entities are described by the global vertex indices of the mesh.
        These entity indices are using the original input ordering.

        Returns:
            Entities owned by the process (and their local entity-to-vertex
            indices) and the corresponding values.
        """
        return _cpp.io.distribute_entity_data(
            mesh.topology._cpp_object,
            mesh.geometry.input_global_indices,
            mesh.geometry.index_map().size_global,
            mesh.geometry.cmap.create_dof_layout(),
            mesh.geometry.dofmap,
            entity_dim,
            entities,
            values,
        )