# Copyright (C) 2018-2021 Michal Habera, Garth N. Wells and Jørgen S. Dokken
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Methods for geometric searches and operations."""

from __future__ import annotations

import typing

import numpy as np
import numpy.typing as npt

if typing.TYPE_CHECKING:
    from dolfinx.cpp.graph import AdjacencyList_int32
    from dolfinx.mesh import Mesh


from dolfinx import cpp as _cpp

__all__ = [
    "BoundingBoxTree",
    "bb_tree",
    "compute_colliding_cells",
    "squared_distance",
    "compute_closest_entity",
    "compute_collisions_trees",
    "compute_collisions_points",
    "compute_distance_gjk",
    "create_midpoint_tree",
    "PointOwnershipData",
]


class PointOwnershipData:
    """Convenience class for storing data related to the ownership of points."""

    _cpp_object: typing.Union[
        _cpp.geometry.PointOwnershipData_float32, _cpp.geometry.PointOwnershipData_float64
    ]

    def __init__(self, ownership_data):
        """Wrap a C++ PointOwnershipData."""
        self._cpp_object = ownership_data

    def src_owner(self) -> npt.NDArray[np.int32]:
        """Ranks owning each point sent into ownership determination for current process"""
        return self._cpp_object.src_owner

    def dest_owner(self) -> npt.NDArray[np.int32]:
        """Ranks that sent `dest_points` to current process"""
        return self._cpp_object.dest_owners

    def dest_points(self) -> npt.NDArray[np.floating]:
        """Points owned by current rank"""
        return self._cpp_object.dest_points

    def dest_cells(self) -> npt.NDArray[np.int32]:
        """Cell indices (local to process) where each entry of `dest_points` is located"""
        return self._cpp_object.dest_cells


class BoundingBoxTree:
    """Bounding box trees used in collision detection."""

    _cpp_object: typing.Union[
        _cpp.geometry.BoundingBoxTree_float32, _cpp.geometry.BoundingBoxTree_float64
    ]

    def __init__(self, tree):
        """Wrap a C++ BoundingBoxTree.

        Note:
            This initializer should not be used in user code. Use
                ``bb_tree``.

        """
        self._cpp_object = tree

    @property
    def num_bboxes(self) -> int:
        """Number of bounding boxes."""
        return self._cpp_object.num_bboxes

    def get_bbox(self, i) -> npt.NDArray[np.floating]:
        """Get lower and upper corners of the ith bounding box.

        Args:
            i: Index of the box.

        Returns:
            The 'lower' and 'upper' points of the bounding box.
            Shape is ``(2, 3)``,

        """
        return self._cpp_object.get_bbox(i)

    def create_global_tree(self, comm) -> BoundingBoxTree:
        return BoundingBoxTree(self._cpp_object.create_global_tree(comm))


def bb_tree(
    mesh: Mesh,
    dim: int,
    entities: typing.Optional[npt.NDArray[np.int32]] = None,
    padding: float = 0.0,
) -> BoundingBoxTree:
    """Create a bounding box tree for use in collision detection.

    Args:
        mesh: The mesh.
        dim: Dimension of the mesh entities to build bounding box for.
        entities: List of entity indices (local to process). If not
            supplied, all owned and ghosted entities are used.
        padding: Padding for each bounding box.

    Returns:
        Bounding box tree.

    """
    map = mesh.topology.index_map(dim)
    if map is None:
        raise RuntimeError(f"Mesh entities of dimension {dim} have not been created.")
    if entities is None:
        entities = np.arange(map.size_local + map.num_ghosts, dtype=np.int32)

    dtype = mesh.geometry.x.dtype
    if np.issubdtype(dtype, np.float32):
        return BoundingBoxTree(
            _cpp.geometry.BoundingBoxTree_float32(mesh._cpp_object, dim, entities, padding)
        )
    elif np.issubdtype(dtype, np.float64):
        return BoundingBoxTree(
            _cpp.geometry.BoundingBoxTree_float64(mesh._cpp_object, dim, entities, padding)
        )
    else:
        raise NotImplementedError(f"Type {dtype} not supported.")


def compute_collisions_trees(
    tree0: BoundingBoxTree, tree1: BoundingBoxTree
) -> npt.NDArray[np.int32]:
    """Compute all collisions between two bounding box trees.

    Args:
        tree0: First bounding box tree.
        tree1: Second bounding box tree.

    Returns:
        List of pairs of intersecting box indices from each tree. Shape
        is ``(num_collisions, 2)``.

    """
    return _cpp.geometry.compute_collisions_trees(tree0._cpp_object, tree1._cpp_object)


def compute_collisions_points(
    tree: BoundingBoxTree, x: npt.NDArray[np.floating]
) -> _cpp.graph.AdjacencyList_int32:
    """Compute collisions between points and leaf bounding boxes.

    Bounding boxes can overlap, therefore points can collide with more
    than one box.

    Args:
        tree: Bounding box tree.
        x: Points (``shape=(num_points, 3)``).

    Returns:
       For each point, the bounding box leaves that collide with the
       point.

    """
    return _cpp.geometry.compute_collisions_points(tree._cpp_object, x)


def compute_closest_entity(
    tree: BoundingBoxTree,
    midpoint_tree: BoundingBoxTree,
    mesh: Mesh,
    points: npt.NDArray[np.floating],
) -> npt.NDArray[np.int32]:
    """Compute closest mesh entity to a point.

    Args:
        tree: bounding box tree for the entities.
        midpoint_tree: A bounding box tree with the midpoints of all
            the mesh entities. This is used to accelerate the search.
        mesh: The mesh.
        points: The points to check for collision, ``shape=(num_points,3)``.

    Returns:
        Mesh entity index for each point in ``points``. Returns -1 for a
        point if the bounding box tree is empty.

    """
    return _cpp.geometry.compute_closest_entity(
        tree._cpp_object, midpoint_tree._cpp_object, mesh._cpp_object, points
    )


def create_midpoint_tree(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32]) -> BoundingBoxTree:
    """Create a bounding box tree for the midpoints of a subset of entities.

    Args:
        mesh: The mesh.
        dim: Topological dimension of the entities.
        entities: Indices of mesh entities to include.

    Returns:
        Bounding box tree for midpoints of cell entities.

    """
    return BoundingBoxTree(_cpp.geometry.create_midpoint_tree(mesh._cpp_object, dim, entities))


def compute_colliding_cells(
    mesh: Mesh, candidates: AdjacencyList_int32, x: npt.NDArray[np.floating]
):
    """From a mesh, find which cells collide with a set of points.

    Args:
        mesh: The mesh.
        candidate_cells: Adjacency list of candidate colliding cells for
            the ith point in ``x``.
        points: The points to check for collision ``shape=(num_points, 3)``,

    Returns:
        Adjacency list where the ith node is the list of entities that
        collide with the ith point.

    """
    return _cpp.geometry.compute_colliding_cells(mesh._cpp_object, candidates, x)


def squared_distance(mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating]):
    """Compute the squared distance between a point and a mesh entity.

    The distance is computed between the ith input points and the ith
    input entity.

    Args:
        mesh: Mesh containing the entities.
        dim: Topological dimension of the mesh entities.
        entities: Indices of the mesh entities (local to process).
        points: Points to compute the shortest distance from
            (``shape=(num_points, 3)``).

    Returns:
        Squared shortest distance from ``points[i]`` to ``entities[i]``.

    """
    return _cpp.geometry.squared_distance(mesh._cpp_object, dim, entities, points)


def compute_distance_gjk(
    p: npt.NDArray[np.floating], q: npt.NDArray[np.floating]
) -> npt.NDArray[np.floating]:
    """Compute the distance between two convex bodies p and q, each defined by a set of points.

    Uses the Gilbert-Johnson-Keerthi (GJK) distance algorithm.

    Args:
        p: Body 1 list of points (``shape=(num_points, gdim)``).
        q: Body 2 list of points (``shape=(num_points, gdim)``).

    Returns:
        Shortest vector between the two bodies.

    """
    assert p.dtype == q.dtype
    if np.issubdtype(p.dtype, np.float32):
        return _cpp.geometry.compute_distance_gjk_float32(p, q)
    elif np.issubdtype(p.dtype, np.float64):
        return _cpp.geometry.compute_distance_gjk_float64(p, q)
    raise RuntimeError("Invalid dtype in compute_distance_gjk")