from __future__ import annotations

import itertools
from collections.abc import Sequence
from typing import TYPE_CHECKING, cast, overload

import numpy as np
import pyvista as pv
from scipy.spatial import KDTree


if TYPE_CHECKING:
    from typing import Any, Literal, Optional  # pragma: no cover

    from numpy.typing import ArrayLike, NDArray  # pragma: no cover


def average_points(mesh: pv.PolyData, tolerance: float = 0.0) -> pv.PolyData:
    """
    Average duplicate points in this mesh.

    Parameters
    ----------
    mesh : pyvista.PolyData
        Mesh to average points from.
    tolerance : float, default 0.0
        Specify a tolerance to use when comparing points. Points within this tolerance
        will be averaged.

    Returns
    -------
    pyvista.PolyData
        Mesh with averaged points.

    """

    def decimate(cell: ArrayLike, close: bool) -> NDArray:
        cell = np.asanyarray(cell)
        cell = cell[np.insert(np.diff(cell), 0, 1) != 0]

        return cell[:-1] if close and cell[0] == cell[-1] else cell

    points = mesh.points
    groups, group_map = [], {}

    for i, j in KDTree(points).query_pairs(tolerance):
        igrp = group_map[i] if i in group_map else -1
        jgrp = group_map[j] if j in group_map else -1

        if igrp >= 0 and jgrp < 0:
            groups[igrp].append(j)
            group_map[j] = igrp

        elif igrp < 0 and jgrp >= 0:
            groups[jgrp].append(i)
            group_map[i] = jgrp

        elif igrp >= 0 and jgrp >= 0:
            if igrp != jgrp:
                group_map.update({k: igrp for k in groups[jgrp]})
                groups[igrp] += groups[jgrp]
                groups[jgrp] = []

        else:
            gid = len(groups)
            groups.append([i, j])
            group_map[i] = gid
            group_map[j] = gid

    point_map = np.arange(mesh.n_points)
    new_points = points.copy()

    for group in groups:
        if not group:
            continue

        point_map[group] = group[0]
        new_points[group[0]] = points[group].mean(axis=0)

    if mesh.n_faces_strict:
        irregular_faces = [
            decimate(point_map[face], close=True) for face in mesh.irregular_faces
        ]
        faces = [face for face in irregular_faces if face.size > 2]
        new_mesh = pv.PolyData().from_irregular_faces(new_points, faces)
        new_mesh.cell_data["vtkOriginalCellIds"] = [
            i for i, face in enumerate(irregular_faces) if face.size > 2
        ]

    else:
        new_mesh = pv.PolyData(new_points)

    if mesh.n_lines:
        raise NotImplementedError()

    return cast(pv.PolyData, new_mesh.clean())


def decimate_rdp(mesh: pv.PolyData, tolerance: float = 1.0e-8) -> pv.PolyData:
    """
    Decimate polylines and/or polygons in a polydata.

    Parameters
    ----------
    mesh : pyvista.PolyData
        Polydata to decimate.
    tolerance : scalar, default 1.0e-8
        Tolerance for the Ramer-Douglas-Peucker algorithm.

    Returns
    -------
    pyvista.PolyData
        Decimated polydata.

    """

    def decimate(points: ArrayLike) -> NDArray:
        """Ramer-Douglas-Packer algorithm."""
        points = np.asanyarray(points)
        u = points[-1] - points[0]
        un = np.linalg.norm(u)
        dist = (
            np.linalg.norm(np.cross(u, points[0] - points), axis=1) / un
            if un > 0.0
            else np.linalg.norm(points - points[0], axis=1)
        )
        imax = dist.argmax()

        if dist[imax] > tolerance:
            res1 = decimate(points[: imax + 1])
            res2 = decimate(points[imax:])

            return np.vstack((res1[:-1], res2))

        else:
            return np.vstack((points[0], points[-1]))

    lines = []
    points = []

    for cell in mesh.cell:
        if cell.type.name in {"LINE", "POLY_LINE"}:
            points_ = mesh.points[cell.point_ids]

            if cell.type.name == "POLY_LINE":
                points_ = decimate(points_)

            lines += [len(points_), *(np.arange(len(points_)) + len(points))]
            points += points_.tolist()

    return cast(pv.PolyData, pv.PolyData(points, lines=lines).clean())


@overload
def extract_boundary_polygons(
    mesh: pv.DataSet,
    fill: Literal[True],
    with_holes: Literal[True],
) -> tuple[pv.UnstructuredGrid, ...] | None: ...


@overload
def extract_boundary_polygons(
    mesh: pv.DataSet,
    fill: Literal[False],
    with_holes: Literal[True],
) -> tuple[list[pv.PolyData], ...] | None: ...


@overload
def extract_boundary_polygons(
    mesh: pv.DataSet,
    fill: Literal[True],
    with_holes: Literal[False],
) -> tuple[pv.PolyData, ...] | None: ...


@overload
def extract_boundary_polygons(
    mesh: pv.DataSet,
    fill: Literal[False],
    with_holes: Literal[False],
) -> tuple[pv.PolyData, ...] | None: ...


@overload
def extract_boundary_polygons(
    mesh: pv.DataSet,
    fill: Literal[False],
) -> tuple[pv.PolyData, ...] | None: ...


def extract_boundary_polygons(
    mesh: pv.DataSet,
    fill: bool = False,
    with_holes: bool = False,
) -> (
    tuple[pv.UnstructuredGrid, ...]
    | tuple[pv.PolyData, ...]
    | tuple[list[pv.PolyData], ...]
    | None
):
    """
    Extract boundary edges of a mesh as continuous polylines or polygons.

    Parameters
    ----------
    mesh : pyvista.DataSet
        Mesh to extract boundary edges from.
    fill : bool, default False
        If True, return boundary edges as polygons.
    with_holes : bool, default False
        If True, group holes with their corresponding boundary edges.

    Returns
    -------
    tuple[pyvista.PolyData | pyvista.UnstructuredGrid, ...] | tuple[list[pyvista.PolyData], ...] | None
        Extracted boundary polylines or polygons.

    """
    import shapely

    from .. import Polygon

    if isinstance(mesh, pv.PolyData) and mesh.n_faces_strict == 0 and mesh.n_lines > 0:
        edges = mesh

    else:
        edges = (
            mesh.cast_to_unstructured_grid()
            .clean()
            .extract_feature_edges(
                boundary_edges=True,
                non_manifold_edges=False,
                feature_edges=False,
                manifold_edges=False,
                clear_data=True,
            )
        )

    edges = edges.strip(max_length=mesh.n_points)
    edges = cast(pv.PolyData, edges)

    if edges.n_lines == 0:
        return None

    edges = [
        cast(pv.PolyData, edge.merge_points())
        for edge in split_lines(edges, as_lines=False)
    ]

    # Identify holes
    if with_holes:
        holes = {}
        polygons = [shapely.Polygon(edge.points) for edge in edges]

        for i, j in itertools.permutations(range(len(polygons)), 2):
            if i in holes or j in holes:
                continue

            if polygons[i].area > polygons[j].area and polygons[i].contains(
                polygons[j]
            ):
                holes[j] = i

        # Group boundary edges and holes
        polygons = [[edge] if i not in holes else [] for i, edge in enumerate(edges)]

        for k, v in holes.items():
            polygons[v].append(edges[k])

        polygons = [polygon for polygon in polygons if polygon]
        polygons = (
            [Polygon(polygon[0], polygon[1:]) for polygon in polygons]
            if fill
            else polygons
        )

    else:
        polygons = (
            [
                polygon
                + pv.PolyData().from_regular_faces(
                    polygon.points, np.expand_dims(np.arange(polygon.n_points), axis=0)
                )
                for polygon in edges
            ]
            if fill
            else edges
        )

    return tuple(polygons)


def extract_cell_geometry(
    mesh: pv.DataSet,
    remove_ghost_cells: bool = True,
) -> pv.PolyData:
    """
    Extract the geometry of individual cells.

    Parameters
    ----------
    mesh : pyvista.DataSet
        Mesh to extract cell geometry from.
    remove_ghost_cells : bool, default True
        If True, remove ghost cells.

    Returns
    -------
    pyvista.PolyData
        Extracted cell geometry.

    """

    def get_polydata_from_points_cells(
        points: ArrayLike,
        cells: Sequence[NDArray] | Sequence[Sequence[int]] | Sequence[Any],
        key: Literal["faces", "lines"],
    ) -> pv.PolyData:
        points = np.asanyarray(points)
        cell_ids, cells_, lines_or_faces = [], [], []
        cell_map = {}

        for i, cell in enumerate(cells):
            if len(cell) == 0:
                continue

            for c in cell:
                cell_set = tuple(sorted(set(c)))

                try:
                    idx = cell_map[cell_set]
                    cell_ids[idx].append(i)

                except KeyError:
                    idx = len(cell_map)
                    cell_map[cell_set] = idx
                    cell_ids.append([i])

                    cells_.append(c)
                    lines_or_faces += [len(c), *cells_[idx]]

        tmp = -np.ones((len(cell_ids), len(max(cell_ids, key=len))), dtype=int)
        for i, ids in enumerate(cell_ids):
            tmp[i, : len(ids)] = ids

        poly = (
            pv.PolyData(points, lines=lines_or_faces)
            if key == "lines"
            else pv.PolyData(points, faces=lines_or_faces)
        )
        poly.cell_data["vtkOriginalCellIds"] = tmp

        return poly

    from .. import get_cell_connectivity, get_dimension

    if not remove_ghost_cells and "vtkGhostType" in mesh.cell_data:
        mesh = mesh.copy(deep=False)
        mesh.clear_data()

    ndim = get_dimension(mesh)
    mesh = mesh.cast_to_unstructured_grid()
    celltypes = mesh.celltypes
    connectivity = get_cell_connectivity(mesh)

    if ndim in {1, 2}:
        supported_celltypes = {
            pv.CellType.EMPTY_CELL,
            pv.CellType.LINE,
            pv.CellType.PIXEL,
            pv.CellType.POLYGON,
            pv.CellType.POLY_LINE,
            pv.CellType.QUAD,
            pv.CellType.TRIANGLE,
        }
        unsupported_celltypes = set(celltypes).difference(supported_celltypes)

        if unsupported_celltypes:
            raise NotImplementedError(
                f"cells of type '{pv.CellType(list(unsupported_celltypes)[0]).name}' are not supported yet"
            )

        # Generate edge data
        pixel_ind = np.array([0, 1, 3, 2])
        cell_edges = [
            np.column_stack((cell[:-1], cell[1:]))
            if celltype in {pv.CellType.LINE, pv.CellType.POLY_LINE}
            else np.column_stack(
                (
                    cell[pixel_ind],
                    np.roll(cell[pixel_ind], -1),
                )
            )
            if celltype == pv.CellType.PIXEL
            else np.column_stack((cell, np.roll(cell, -1)))
            if not remove_ghost_cells or celltype != pv.CellType.EMPTY_CELL
            else []
            for cell, celltype in zip(connectivity, celltypes)
        ]

        poly = get_polydata_from_points_cells(mesh.points, cell_edges, "lines")

        # Handle collapsed cells
        if remove_ghost_cells:
            lengths = poly.compute_cell_sizes(
                length=True, area=False, volume=False
            ).cell_data["Length"]
            mask = np.abs(lengths) > 0.0

            if not mask.all():
                lines = poly.lines.reshape((poly.n_lines, 3))[mask]
                tmp = poly.cell_data["vtkOriginalCellIds"][mask]
                tmp = tmp[:, np.ptp(tmp, axis=0) > 0]

                poly = pv.PolyData(poly.points, lines=lines)
                poly.cell_data["vtkOriginalCellIds"] = tmp

    else:
        # Generate face data
        if np.ptp(celltypes) == 0:
            celltype = celltypes[0]
            cell_faces = (
                [
                    [
                        face
                        for v in _celltype_to_faces[celltype].values()
                        for face in cell[v]
                    ]
                    for cell in connectivity
                ]
                if celltype != pv.CellType.POLYHEDRON
                else connectivity
            )

        else:
            cell_faces = []

            for cell, celltype in zip(connectivity, celltypes):
                if celltype == pv.CellType.POLYHEDRON:
                    cell_face = cell

                elif celltype in _celltype_to_faces:
                    cell_face = [
                        face
                        for v in _celltype_to_faces[celltype].values()
                        for face in cell[v]
                    ]

                elif celltype == pv.CellType.EMPTY_CELL:
                    cell_face = []

                else:
                    raise NotImplementedError(
                        f"cells of type '{celltype.name}' are not supported yet"
                    )

                cell_faces.append(cell_face)

        poly = get_polydata_from_points_cells(mesh.points, cell_faces, "faces")

        # Handle collapsed cells
        if remove_ghost_cells:
            areas = poly.compute_cell_sizes(
                length=False, area=True, volume=False
            ).cell_data["Area"]
            mask = np.abs(areas) > 0.0

            if not mask.all():
                faces = [
                    face for face, mask_ in zip(poly.irregular_faces, mask) if mask_
                ]
                tmp = poly.cell_data["vtkOriginalCellIds"][mask]
                tmp = tmp[:, np.ptp(tmp, axis=0) > 0]

                poly = pv.PolyData().from_irregular_faces(poly.points, faces)
                poly.cell_data["vtkOriginalCellIds"] = tmp

    return poly


def extract_cells(
    mesh: pv.DataSet,
    ind: ArrayLike,
    invert: bool = False,
    progress_bar: bool = False,
) -> pv.UnstructuredGrid:
    """
    Get a subset of the grid.

    Parameters
    ----------
    ind : ArrayLike
        Indices of cells to extract.
    invert : bool, default False
        If True, invert the selection.
    progress_bar : bool, default False
        If True, display a progress bar.

    Returns
    -------
    pyvista.UnstructuredGrid
        Extracted cells.

    Note
    ----
    This function wraps `pyvista.DataSet.extract_cells()` with consistent handling of
    ghost cells across different versions of VTK.

    """
    ind = np.asanyarray(ind)
    ghost_cells = (
        mesh.cell_data.pop("vtkGhostType") if "vtkGhostType" in mesh.cell_data else None
    )

    try:
        cells = mesh.extract_cells(ind, invert=invert, progress_bar=progress_bar)

        if ghost_cells is not None:
            cells.cell_data["vtkGhostType"] = ghost_cells[
                cells.cell_data["vtkOriginalCellIds"]
            ]

    finally:
        if ghost_cells is not None:
            mesh.cell_data["vtkGhostType"] = ghost_cells

    return cast(pv.UnstructuredGrid, cells)


def extract_cells_by_dimension(
    mesh: pv.DataSet,
    ndim: Optional[int] = None,
    method: Literal["lower", "upper"] = "upper",
    keep_empty_cells: bool = False,
) -> pv.UnstructuredGrid:
    """
    Extract cells by a specified dimension.

    Parameters
    ----------
    mesh : pyvista.DataSet
        Mesh to extract cells from.
    ndim : int, optional
        Dimension to be used for extraction. If None, the dimension of *mesh* is used.
    method : {'lower', 'upper'}, default 'upper'
        Set the extraction method. 'lower' will extract cells of dimension lower than
        *ndim*. 'upper' will extract cells of dimension larger than *ndim*.
    keep_empty_cells : bool, default False
        If True, keep empty cells in the output mesh.

    Returns
    -------
    pyvista.UnstructuredGrid
        Mesh with extracted cells.

    """
    from ._properties import _dimension_map
    from .. import get_dimension

    mesh = mesh.cast_to_unstructured_grid()
    ndim = ndim if ndim is not None else get_dimension(mesh)

    if method == "upper":
        mask = _dimension_map[mesh.celltypes] >= ndim

    elif method == "lower":
        mask = _dimension_map[mesh.celltypes] <= ndim

    else:
        raise ValueError(f"invalid method '{method}' (expected 'lower' or 'upper')")

    if keep_empty_cells:
        mask |= mesh.celltypes == pv.CellType.EMPTY_CELL

    if not mask.all():
        mesh = extract_cells(mesh, mask)

    return mesh


def fuse_cells(
    mesh: pv.DataSet, ind: Sequence[int] | Sequence[Sequence[int]]
) -> pv.UnstructuredGrid:
    """
    Fuse connected cells into a single cell.

    Parameters
    ----------
    mesh : pyvista.DataSet
        Mesh to fuse cells from.
    ind : Sequence[int] | Sequence[Sequence[int]]
        Indices or sequence of indices of cells to fuse.

    Returns
    -------
    pyvista.UnstructuredGrid
        Mesh with fused cells.

    """
    from .. import extract_boundary_polygons, get_cell_connectivity, get_dimension

    indices = [ind] if np.ndim(ind[0]) == 0 else ind
    indices = cast(Sequence[Sequence[int]], indices)
    mesh = mesh.cast_to_unstructured_grid()
    connectivity = list(get_cell_connectivity(mesh))
    celltypes = mesh.celltypes.copy()
    mask = np.ones(mesh.n_cells, dtype=bool)

    for ind_ in indices:
        ind_ = np.asanyarray(ind_)
        ind_ = np.flatnonzero(ind_) if ind_.dtype.kind == "b" else ind_
        mesh_ = extract_cells(mesh, ind_)
        mask[ind_[1:]] = False

        if get_dimension(mesh_) == 2:
            poly = extract_boundary_polygons(mesh_, fill=False)

            if poly is None or len(poly) == 0:
                raise ValueError(
                    "could not extract boundary polygons for the selected cells"
                )

            if len(poly) > 1:
                raise ValueError("could not fuse not fully connected cells together")

            # Find original point IDs of polygon
            # Select first instance found for each point
            cell = poly[0].cast_to_unstructured_grid()
            ids = np.array(
                [
                    np.flatnonzero(mask)[0]
                    for mask in (
                        cell.points[:, None] == mesh.points.astype(cell.points.dtype)
                    ).all(axis=-1)
                ]
            )
            mesh_points = mesh.points[ids]
            sorted_ids = ids[
                np.ravel(
                    [
                        np.flatnonzero(
                            (mesh_points.astype(point.dtype) == point).all(axis=1)
                        )
                        for point in cell.points
                    ]
                )
            ]

            # Update connectivity and cell type
            connectivity[ind_[0]] = sorted_ids
            celltypes[ind_[0]] = pv.CellType.POLYGON

            for cell_id in ind_[1:]:
                connectivity[cell_id] = []
                celltypes[cell_id] = pv.CellType.EMPTY_CELL

            # Generate new mesh with fused cells
            cells = [item for cell in connectivity for item in [len(cell), *cell]]

        else:
            raise NotImplementedError("could not fuse cells for non 2D mesh")

    fused_mesh = pv.UnstructuredGrid(cells, celltypes, mesh.points)
    fused_mesh.point_data.update(mesh.point_data)
    fused_mesh.cell_data.update(mesh.cell_data)
    fused_mesh.user_dict.update(mesh.user_dict)

    # Tidy up
    fused_mesh = extract_cells(fused_mesh, mask).clean()

    return cast(pv.UnstructuredGrid, fused_mesh)


def intersect_polyline(
    mesh: pv.DataSet,
    line: pv.PolyData,
    min_length: float = 1.0e-4,
    tolerance: float = 1.0e-8,
    pass_cell_data: bool = True,
    ignore_points_before_entry: bool = False,
    ignore_points_after_exit: bool = False,
) -> pv.PolyData:
    """
    Intersect a polyline with a mesh.

    Parameters
    ----------
    mesh : pyvista.DataSet
        Mesh to intersect with.
    line : pyvista.PolyData
        Polyline to intersect with the mesh.
    min_length : scalar, default 1.0e-4
        Set the minimum length of a line.
    tolerance : float, default 1.0e-8
        The absolute tolerance to use to find cells along the line.
    pass_cell_data : bool, default True
        If True, pass cell data from the line to the output.
    ignore_points_before_entry : bool, default False
        If True, ignore points before the first entry point into the mesh.
    ignore_points_after_exit : bool, default False
        If True, ignore points after the last exit point from the mesh.

    Returns
    -------
    pyvista.PolyData
        Polydata containing the intersection points and cell IDs.

    """
    from .. import get_cell_centers

    line_ = cast(pv.PolyData, line.strip())
    lines = split_lines(line_, as_lines=True)[0]

    # Recenter coordinates around zero to prevent accuracy issues
    center = np.array(mesh.center)
    mesh = mesh.translate(-center)
    lines = lines.translate(-center)

    line_ids, cell_ids, cell = [], [], None
    mesh_entered, mesh_exited = False, False
    points = [lines.points[0]]
    count = 0

    def add_point(point: ArrayLike, line_id: int, cell_id: int) -> None:
        """Add a point to the intersection results."""
        if not np.allclose(points[-1], point, atol=tolerance):
            if (
                not mesh_entered
                or mesh_exited
                or np.linalg.norm(point - points[-1]) >= min_length
            ):
                points.append(point)
                line_ids.append(line_id)
                cell_ids.append(cell_id)

    cell_geometry = extract_cell_geometry(mesh, remove_ghost_cells=True)
    tree = KDTree(get_cell_centers(cell_geometry))

    for lid, (pointa, pointb) in enumerate(zip(lines.points[:-1], lines.points[1:])):
        # Find the first cell intersected by the line
        if cell is None:
            ids = mesh.find_cells_intersecting_line(pointa, pointb, tolerance=tolerance)

            if ids.size == 0:
                add_point(pointb, lid, -1)
                continue

            ids = np.sort(ids)
            id_ = np.linalg.norm(
                get_cell_centers(extract_cells(mesh, ids)) - pointa,
                axis=-1,
            ).argmin()
            cid = ids[id_]
            cell = extract_cells(mesh, cid)

        if not mesh_exited:
            while True:
                # Find cell edges or faces intersected by the line
                intersections = ray_cast(
                    cell,
                    pointa,
                    pointb,
                    tolerance=tolerance,
                )

                # Find the exit, if any
                # Line is fully contained in the cell
                if intersections is None:
                    break

                # Either an entrance or an exit
                elif intersections.n_cells == 1:
                    # It's an entrance if intersection point matches previous point
                    if not mesh_entered or np.allclose(
                        points[-1],
                        intersections.cell_data["IntersectionPoints"][0],
                        atol=tolerance,
                    ):
                        if not mesh_entered:
                            count = len(points)
                            mesh_entered = True
                            add_point(
                                intersections.cell_data["IntersectionPoints"][0],
                                lid,
                                -1,
                            )

                            if count > 1:
                                break

                        else:
                            add_point(
                                intersections.cell_data["IntersectionPoints"][0],
                                lid,
                                int(cid),
                            )
                            break

                    fid = 0

                # The exit is the farthest from last point
                elif intersections.n_cells == 2:
                    dist = np.linalg.norm(
                        intersections.cell_data["IntersectionPoints"] - points[-1],
                        axis=-1,
                    )

                    if not mesh_entered:
                        count = len(points)
                        add_point(
                            intersections.cell_data["IntersectionPoints"][
                                dist.argmin()
                            ],
                            lid,
                            -1,
                        )
                        mesh_entered = True

                    fid = dist.argmax()

                # The line is hitting an edge or a corner
                else:
                    raise ValueError("could not find the exit face")

                add_point(
                    intersections.cell_data["IntersectionPoints"][fid], lid, int(cid)
                )
                exit_face = intersections.get_cell(fid)

                # Check distance to last point
                if (
                    np.linalg.norm(
                        intersections.cell_data["IntersectionPoints"][fid]
                        - lines.points[-1]
                    )
                    < min_length
                ):
                    break

                # Determine the exit cell
                _, id_ = tree.query(exit_face.center)
                cells = cell_geometry.cell_data["vtkOriginalCellIds"][id_]

                if cid not in cells:
                    raise ValueError("could not determine exit cell")

                if -1 in cells:
                    if not ignore_points_after_exit:
                        add_point(pointb, lid, -1)

                    mesh_exited = True
                    break

                cid = [id_ for id_ in cells if id_ != cid][0]
                cell = extract_cells(mesh, cid)

                # Add last point of the polyline if it is inside the exit cell
                if lid == lines.n_lines - 1 and cell.find_containing_cell(pointb) > -1:
                    add_point(pointb, lid, cid)
                    break

        else:
            if ignore_points_after_exit:
                break

            add_point(pointb, lid, -1)

    if ignore_points_before_entry:
        points = points[count:]
        line_ids = line_ids[count:]
        cell_ids = cell_ids[count:]

    polyline = split_lines(pv.MultipleLines(points))[0].compute_cell_sizes(
        length=True, area=False, volume=False
    )
    polyline.cell_data["vtkOriginalCellIds"] = line_ids
    polyline.cell_data["IntersectedCellIds"] = cell_ids

    if pass_cell_data:
        for k, v in line.cell_data.items():
            if k in polyline.cell_data:
                continue

            polyline.cell_data[k] = v[line_ids]

    return cast(pv.PolyData, polyline.translate(center))


def merge(
    dataset: Sequence[pv.StructuredGrid | pv.UnstructuredGrid],
    axis: Optional[int] = None,
    merge_points: bool = True,
) -> pv.StructuredGrid | pv.UnstructuredGrid:
    """
    Merge several meshes.

    Parameters
    ----------
    dataset : Sequence[pyvista.StructuredGrid | pyvista.UnstructuredGrid]
        Meshes to merge together. At least two meshes are required.
    axis : int, optional
        The axis along which two structured grids are merged (if all meshes in *dataset*
        are structured grids).
    merge_points : bool, default True
        If True, merge equivalent points for two unstructured grids.

    Returns
    -------
    pyvista.StructuredGrid | pyvista.UnstructuredGrid
        Merged mesh.

    """
    if len(dataset) < 2:
        return dataset[0]

    if all(isinstance(mesh, pv.StructuredGrid) for mesh in dataset):
        if axis is None:
            raise ValueError("could not merge structured grids with None axis")

        mesh_a = dataset[0]

        for mesh_b in dataset[1:]:
            if axis == 0:
                if not (
                    np.allclose(mesh_a.x[-1], mesh_b.x[0])
                    and np.allclose(mesh_a.y[-1], mesh_b.y[0])
                    and np.allclose(mesh_a.z[-1], mesh_b.z[0])
                ):
                    raise ValueError(
                        "could not merge structured grids with non-matching east and west interfaces"
                    )

                slice_ = (slice(1, None),)

            elif axis == 1:
                if not (
                    np.allclose(mesh_a.x[:, -1], mesh_b.x[:, 0])
                    and np.allclose(mesh_a.y[:, -1], mesh_b.y[:, 0])
                    and np.allclose(mesh_a.z[:, -1], mesh_b.z[:, 0])
                ):
                    raise ValueError(
                        "could not merge structured grids with non-matching north and south interfaces"
                    )

                slice_ = (slice(None), slice(1, None))

            else:
                if not (
                    np.allclose(mesh_a.x[..., -1], mesh_b.x[..., 0])
                    and np.allclose(mesh_a.y[..., -1], mesh_b.y[..., 0])
                    and np.allclose(mesh_a.z[..., -1], mesh_b.z[..., 0])
                ):
                    raise ValueError(
                        "could not merge structured grids with non-matching top and bottom interfaces"
                    )

                slice_ = (slice(None), slice(None), slice(1, None))

            X = np.concatenate((mesh_a.x, mesh_b.x[slice_]), axis=axis)
            Y = np.concatenate((mesh_a.y, mesh_b.y[slice_]), axis=axis)
            Z = np.concatenate((mesh_a.z, mesh_b.z[slice_]), axis=axis)
            mesh = pv.StructuredGrid(X, Y, Z)

            if mesh_a.cell_data:
                shape_a = [max(1, n - 1) for n in mesh_a.dimensions]
                shape_b = [max(1, n - 1) for n in mesh_b.dimensions]
                mesh.cell_data.update(
                    {
                        k: np.concatenate(
                            (
                                v.reshape(shape_a, order="F"),
                                mesh_b.cell_data[k].reshape(shape_b, order="F"),
                            ),
                            axis=axis,
                        ).ravel(order="F")
                        for k, v in mesh_a.cell_data.items()
                        if k in mesh_b.cell_data
                    }
                )

            mesh_a = mesh

    else:
        mesh = pv.merge(dataset, merge_points=merge_points)

    return mesh


def merge_lines(
    lines: Sequence[pv.PolyData],
    as_lines: bool = True,
) -> pv.PolyData:
    """
    Merge line(s) or polyline(s) into a polydata.

    Parameters
    ----------
    lines : Sequence[pyvista.PolyData]
        List of line(s) or polyline(s) to merge.
    as_lines : bool, default True
        If True, return merged line(s) or polyline(s) as line(s).

    Returns
    -------
    pyvista.PolyData
        Polydata with merged line(s) or polyline(s).

    Note
    ----
    Preserve ordering compared to pyvista.merge().

    """
    # Find common point and cell data keys
    point_data_keys = set(lines[0].point_data)
    cell_data_keys = set(lines[0].cell_data)

    for lines_ in lines[1:]:
        point_data_keys = point_data_keys.intersection(lines_.point_data)
        cell_data_keys = cell_data_keys.intersection(lines_.cell_data)

    point_data = {k: [] for k in point_data_keys}
    cell_data = {k: [] for k in cell_data_keys}

    # Loop over lines
    points, cells, offset = [], [], 0

    for lines_ in lines:
        for line in split_lines(lines_, as_lines=False):
            points.append(line.points)
            ids = np.arange(line.n_points) + offset
            cells += (
                np.insert(np.column_stack((ids[:-1], ids[1:])), 0, 2, axis=-1)
                .ravel()
                .tolist()
                if as_lines
                else [line.n_points, *ids]
            )
            offset += line.n_points

            for k, v in point_data.items():
                v.append(line.point_data[k])

            for k, v in cell_data.items():
                v.append(
                    (
                        np.full(ids.size - 1, line.cell_data[k])
                        if np.ndim(line.cell_data[k]) == 0
                        or line.cell_data[k].dtype.kind == "U"
                        else np.tile(line.cell_data[k], (ids.size - 1, 1))
                    )
                    if as_lines
                    else line.cell_data[k]
                )

    mesh = pv.PolyData(np.concatenate(points), lines=cells)

    for k, v in point_data.items():
        mesh.point_data[k] = np.concatenate(v) if v[0].ndim == 1 else np.vstack(v)

    for k, v in cell_data.items():
        mesh.cell_data[k] = np.concatenate(v) if v[0].ndim == 1 else np.vstack(v)

    return cast(pv.PolyData, mesh.merge_points())


def offset_polygon(
    mesh_or_points: pv.PolyData | ArrayLike,
    distance: float,
) -> pv.PolyData:
    """
    Offset a polygon by a specified distance.

    Parameters
    ----------
    mesh_or_points : pyvista.PolyData | ArrayLike
        Polygon to offset.
    distance : scalar
        Distance to offset the polygon.

    Returns
    -------
    pyvista.PolyData
        Offset polygon.

    """
    if not isinstance(mesh_or_points, pv.PolyData):
        mesh_or_points = np.asanyarray(mesh_or_points)
        n_points = len(mesh_or_points)
        faces = [n_points, *np.arange(n_points)]
        mesh = pv.PolyData(mesh_or_points, faces=faces)

    else:
        mesh = mesh_or_points

    if not mesh.n_faces_strict:
        raise ValueError("could not offset polygon with zero polygon")

    if distance > 0.0:
        fac = 1.0

    elif distance < 0.0:
        fac = -1.0
        distance *= -1.0

    else:
        return mesh.copy()

    # Loop over faces
    faces = []
    points_ = []

    for face in mesh.irregular_faces:
        # Extract points
        points = mesh.points[face]
        mask = np.ptp(points, axis=0) == 0.0

        if mask.sum() != 1:
            raise ValueError("could not offset non-planar polygon")

        else:
            axis = np.flatnonzero(mask)[0]

        # Simple polygon offset algorithm
        # Vectorized version of C# code
        # <https://stackoverflow.com/a/73061541/9729313>
        points = points[:, ~mask]
        points = np.vstack(
            (
                points[-1],
                points,
                points[0],
            ),
        )

        x, y = points.T
        signed_area = (x[:-1] * y[1:] - x[1:] * y[:-1]).sum()

        vn = points[2:] - points[1:-1]
        vn /= np.linalg.norm(vn, axis=1)[:, np.newaxis]

        vp = points[1:-1] - points[:-2]
        vp /= np.linalg.norm(vp, axis=1)[:, np.newaxis]

        vb = (vn + vp) * fac * np.sign(signed_area)
        vb[:, 0] *= -1.0
        vb /= np.linalg.norm(vb, axis=1)[:, np.newaxis]

        dist = distance / (0.5 * (1.0 + (vn * vp).sum(axis=1))) ** 0.5
        points = points[1:-1] + dist[:, np.newaxis] * vb[:, [1, 0]]

        faces += [len(points), *(np.arange(len(points)) + len(points_))]
        points_ += np.insert(points, axis, 0.0, axis=1).tolist()

    return pv.PolyData(points_, faces=faces)


def ray_cast(
    mesh: pv.DataSet,
    pointa: ArrayLike,
    pointb: ArrayLike,
    tolerance: float = 1.0e-8,
    max_angle: Optional[float] = None,
) -> pv.PolyData | None:
    """
    Perform a single ray casting calculation.

    Parameters
    ----------
    mesh : pyvista.DataSet
        The input mesh to perform ray casting on.
    pointa : ArrayLike
        Length 3 coordinate of the start of the ray.
    pointb : ArrayLike
        Length 3 coordinate of the end of the ray.
    tolerance : float, default 1.0e-8
        The absolute tolerance to use to find cells along the ray.
    max_angle : float, optional
        The maximum angle between face normals and the ray direction. Ignored if *mesh* is 2-dimensional.

    Returns
    -------
    pyvista.PolyData | None
        A polydata with the intersected cells and points. None if no intersections.

    """
    from .. import get_cell_centers

    if isinstance(mesh, pv.PolyData):
        if mesh.n_faces_strict and mesh.n_lines:
            raise ValueError(
                "could not ray cast on a polydata with both faces and lines"
            )

        ids_ = None

    else:
        mesh = extract_cell_geometry(mesh, remove_ghost_cells=False)
        ids_ = mesh.cell_data["vtkOriginalCellIds"]

    pointa = np.asanyarray(pointa)
    pointb = np.asanyarray(pointb)

    # Find cells intersected by the line
    ids = mesh.find_cells_intersecting_line(pointa, pointb, tolerance=tolerance)

    if ids.size == 0:
        return None

    ids = np.sort(ids)
    dvec = (pointa - pointb) / np.linalg.norm(pointa - pointb)

    # Filter faces based on angle with line direction
    if mesh.n_faces_strict:
        max_angle = max_angle if max_angle is not None else 90.0 - tolerance
        normals = mesh.compute_normals(
            cell_normals=True, point_normals=False
        ).cell_data["Normals"]
        angles = np.rad2deg(np.arccos(np.abs(normals[ids] @ dvec)))
        ids = ids[angles < max_angle]

        if ids.size == 0:
            return None

    # Calculate intersection points
    cells = extract_cells(mesh, ids).extract_geometry()

    if mesh.n_faces_strict:
        centers = get_cell_centers(cells)
        intersection = pointa + dvec * np.expand_dims(
            ((centers - pointa) * normals[ids]).sum(axis=1)
            / (dvec * normals[ids]).sum(axis=1),
            axis=1,
        )

    else:
        points = np.array([edge.points for edge in split_lines(cells, as_lines=True)])
        vecs = np.diff(points, axis=1).squeeze()
        vecs = np.atleast_2d(vecs)
        vecs /= np.linalg.norm(vecs, axis=1)[:, None]
        cross = np.cross(vecs, dvec)
        denom = np.linalg.norm(cross, axis=1)
        t = (np.cross((points[:, 0] - pointa), dvec) * cross).sum(axis=1) / denom**2
        intersection = points[:, 0] - t[:, None] * vecs

    # Add data
    cells.clear_data()
    cells.cell_data["IntersectionPoints"] = intersection
    cells.cell_data["vtkOriginalCellIds"] = ids_[ids] if ids_ is not None else ids

    return cells


def reconstruct_line(
    mesh_or_points: pv.DataSet | ArrayLike,
    start: int = 0,
    close: bool = False,
    tolerance: float = 1.0e-8,
) -> pv.PolyData:
    """
    Reconstruct a line from the points in this dataset.

    Parameters
    ----------
    mesh_or_points : pyvista.DataSet | ArrayLike
        Mesh from which points to reconstruct a line.
    start : int, default 0
        Index of point to use as starting point for 2-opt algorithm.
    close : bool, default False
        If True, the ending point is the starting point.
    tolerance : scalar, default 1.0e-8
        Tolerance for the 2-opt algorithm.

    Returns
    -------
    pyvista.PolyData
        Reconstructed line.

    """
    if isinstance(mesh_or_points, pv.PolyData):
        points = mesh_or_points.points

    else:
        points = np.asanyarray(mesh_or_points)

    if not (points.ndim == 2 and points.shape[1] in {2, 3}):
        raise ValueError(
            f"could not reconstruct polyline from {points.shape[1]}D points"
        )

    def path_length(path):
        if close:
            path = np.append(path, path[0])

        return np.linalg.norm(np.diff(points[path], axis=0), axis=1).sum()

    n = len(points)
    shortest_path = np.roll(np.arange(n), -start)
    shortest_length = path_length(shortest_path)
    path = shortest_path.copy()

    while True:
        best_length = shortest_length

        for first in range(1, n - 2):
            for last in range(first + 2, n + 1):
                path[first:last] = np.flip(path[first:last])
                length = path_length(path)

                if length < shortest_length:
                    shortest_path[:] = path
                    shortest_length = length

                else:
                    path[first:last] = np.flip(
                        path[first:last]
                    )  # reset path to current shortest path

        if shortest_length > (1.0 - tolerance) * best_length:
            break

    points = points[shortest_path]
    points = (
        points
        if points.shape[1] == 3
        else np.column_stack((points, np.zeros(len(points))))
    )

    return pv.lines_from_points(points, close=close)


def remap_categorical_data(
    mesh: pv.DataSet,
    key: str,
    mapping: dict[str | int, int],
    preference: Literal["cell", "point"] = "cell",
    inplace: bool = False,
) -> pv.DataSet | None:
    """
    Remap categorical cell or point data.

    Parameters
    ----------
    mesh : pyvista.DataSet
        Mesh with categorical data to remap.
    key : str
        Name of the categorical data to remap.
    mapping : dict
        Mapping of old to new values.
    preference : {'cell', 'point'}, default 'cell'
        Determine whether to remap cell or point data.
    inplace : bool, default False
        If True, overwrite the original mesh.

    Returns
    -------
    pyvista.DataSet | None
        Mesh with remapped categorical data.

    """
    if preference == "cell":
        data = mesh.cell_data[key]

    elif preference == "point":
        data = mesh.point_data[key]

    else:
        raise ValueError(f"invalid preference '{preference}'")

    if data.dtype.kind != "i":
        raise ValueError(f"could not remap non-categorical '{preference}' data '{key}'")

    try:
        data_labels = dict(mesh.user_dict[key])

    except KeyError:
        data_labels = {}

    if not inplace:
        mesh = mesh.copy()

    remapped_data = data.copy()
    data_labels_map = {v: k for k, v in data_labels.items()}
    unused_labels = set(list(data_labels))

    for k, v in mapping.items():
        if isinstance(k, str):
            try:
                vid = data_labels[k]

            except KeyError:
                raise ValueError(f"could not map unknown key '{k}'")

        else:
            vid = k

        mask = data == vid

        if mask.any():
            remapped_data[mask] = v

            try:
                key_ = k if isinstance(k, str) else data_labels_map[vid]
                data_labels[key_] = v
                unused_labels.remove(key_)

            except KeyError:
                pass

    if preference == "cell":
        mesh.cell_data[key] = remapped_data

    else:
        mesh.point_data[key] = remapped_data

    if data_labels:
        for k in unused_labels:
            data_labels.pop(k, None)

        mesh.user_dict[key] = dict(sorted(data_labels.items(), key=lambda x: x[1]))

    if not inplace:
        return mesh


def split_lines(mesh: pv.PolyData, as_lines: bool = True) -> Sequence[pv.PolyData]:
    """
    Split line(s) or polyline(s) into multiple line(s) or polyline(s).

    Parameters
    ----------
    mesh : pyvista.PolyData
        Mesh with line(s) or polyline(s) to split.
    as_lines : bool, default True
        If True, return split line(s) or polyline(s) as line(s).

    Returns
    -------
    Sequence[pyvista.PolyData]
        Split line(s) or polyline(s).

    """
    from pyvista.core.cell import _get_irregular_cells

    mesh = cast(
        pv.PolyData,
        mesh.extract_cells_by_type((pv.CellType.LINE, pv.CellType.POLY_LINE)),
    )
    lines = []

    for i, line_ids in enumerate(_get_irregular_cells(mesh.GetLines())):
        lines_ = (
            np.insert(
                np.column_stack(
                    (np.arange(line_ids.size - 1), np.arange(1, line_ids.size))
                ),
                0,
                2,
                axis=-1,
            ).ravel()
            if as_lines
            else [line_ids.size, *np.arange(line_ids.size)]
        )
        line = pv.PolyData(
            mesh.points[line_ids],
            lines=lines_,
        )

        for k, v in mesh.point_data.items():
            line.point_data[k] = v[line_ids]

        for k, v in mesh.cell_data.items():
            line.cell_data[k] = (
                np.full(line.n_cells, v[i])
                if np.ndim(v[i]) == 0
                else np.tile(v[i], (line.n_cells, 1)).copy()
            )

        lines.append(line)

    return lines


def quadraticize(mesh: pv.UnstructuredGrid) -> pv.UnstructuredGrid:
    """
    Convert linear mesh to quadratic mesh.

    Parameters
    ----------
    mesh : pyvista.UnstructuredGrid
        Mesh with linear cells.

    Returns
    -------
    pyvista.UnstructuredGrid
        Mesh with quadratic cells.

    """
    n_points = mesh.n_points

    cells = []
    celltypes = []
    quad_points = []

    for cell in mesh.cell:
        if cell.type.name not in {"TRIANGLE", "QUAD"}:
            raise NotImplementedError()

        celltype = f"QUADRATIC_{cell.type.name}"
        new_points = 0.5 * (cell.points + np.roll(cell.points, -1, axis=0))
        n_new_points = len(new_points)
        new_points_ids = np.arange(n_new_points) + n_points

        celltypes.append(int(pv.CellType[celltype]))
        cell_ = cell.point_ids + new_points_ids.tolist()
        cells += [len(cell_), *cell_]
        quad_points.append(new_points)
        n_points += n_new_points

    quad_points = np.concatenate(quad_points)
    points = np.vstack((mesh.points, quad_points))

    return pv.UnstructuredGrid(cells, celltypes, points)


_celltype_to_faces = {
    pv.CellType.TETRA: {
        "TRIANGLE": np.array([[1, 2, 3], [0, 3, 2], [0, 1, 3], [0, 2, 1]]),
    },
    pv.CellType.PYRAMID: {
        "QUAD": np.array([[0, 3, 2, 1]]),
        "TRIANGLE": np.array([[0, 1, 4], [1, 2, 4], [2, 3, 4], [3, 0, 4]]),
    },
    pv.CellType.WEDGE: {
        "TRIANGLE": np.array([[0, 2, 1], [3, 4, 5]]),
        "QUAD": np.array([[0, 1, 4, 3], [1, 2, 5, 4], [0, 3, 5, 2]]),
    },
    pv.CellType.HEXAHEDRON: {
        "QUAD": np.array(
            [
                [0, 3, 2, 1],
                [4, 5, 6, 7],
                [0, 1, 5, 4],
                [1, 2, 6, 5],
                [2, 3, 7, 6],
                [0, 4, 7, 3],
            ]
        ),
    },
    pv.CellType.VOXEL: {
        "QUAD": np.array(
            [
                [0, 2, 3, 1],
                [4, 5, 7, 6],
                [0, 1, 5, 4],
                [1, 3, 7, 5],
                [3, 2, 6, 7],
                [0, 4, 6, 2],
            ]
        ),
    },
}
