"""
Ray queries using the embreex package with the
API wrapped to match our native raytracer.
"""

import numpy as np

# `pip install embreex` installs from wheels
from embreex import rtcore_scene
from embreex.mesh_construction import TriangleMesh

from .. import caching, intersections, util
from ..constants import log_time
from ..typed import ArrayLike, Integer
from .ray_util import contains_points

# embree operates on float32 values
_embree_dtype = np.float32

# when calculating multiple hits we offset the hit to
# advance the ray past the plane of the triangle we hit
# hardcode offset for rays larger than our resolution:
#   np.finfo(_embree_dtype).resolution * 10
_ray_offset = 1e-5


class RayMeshIntersector:
    def __init__(self, geometry, scale_to_box: bool = True):
        """
        Do ray- mesh queries.

        Parameters
        -------------
        geometry : Trimesh object
          Mesh to do ray tests on
        scale_to_box : bool
          If true, will scale mesh to approximate
          unit cube to avoid problems with extreme
          large or small meshes.
        """
        self.mesh = geometry
        self._scale_to_box = scale_to_box
        self._cache = caching.Cache(id_function=self.mesh.__hash__)

    @property
    def _scale(self):
        """
        Scaling factor for precision.
        """
        if self._scale_to_box:
            # scale vertices to approximately a cube to help with
            # numerical issues at very large/small scales
            scale = 100.0 / self.mesh.scale
        else:
            scale = 1.0
        return scale

    @caching.cache_decorator
    def _scene(self):
        """
        A cached version of the embreex scene.
        """
        return _EmbreeWrap(
            vertices=self.mesh.vertices, faces=self.mesh.faces, scale=self._scale
        )

    def intersects_location(
        self,
        ray_origins: ArrayLike,
        ray_directions: ArrayLike,
        multiple_hits: bool = True,
    ):
        """
        Return the location of where a ray hits a surface.

        Parameters
        ----------
        ray_origins : (n, 3) float
          Origins of rays
        ray_directions : (n, 3) float
          Direction (vector) of rays

        Returns
        ---------
        locations : (m) sequence of (p, 3) float
          Intersection points
        index_ray : (m,) int
          Indexes of ray
        index_tri : (m,) int
          Indexes of mesh.faces
        """
        (index_tri, index_ray, locations) = self.intersects_id(
            ray_origins=ray_origins,
            ray_directions=ray_directions,
            multiple_hits=multiple_hits,
            return_locations=True,
        )

        return locations, index_ray, index_tri

    @log_time
    def intersects_id(
        self,
        ray_origins: ArrayLike,
        ray_directions: ArrayLike,
        multiple_hits: bool = True,
        max_hits: Integer = 100,
        return_locations: bool = False,
    ):
        """
        Find the triangles hit by a list of rays, including
        optionally multiple hits along a single ray.


        Parameters
        ----------
        ray_origins : (n, 3) float
          Origins of rays
        ray_directions : (n, 3) float
          Direction (vector) of rays
        multiple_hits : bool
          If True will return every hit along the ray
          If False will only return first hit
        max_hits : int
          Maximum number of hits per ray
        return_locations : bool
          Should we return hit locations or not

        Returns
        ---------
        index_tri : (m,) int
          Indexes of mesh.faces
        index_ray : (m,) int
          Indexes of ray
        locations : (m) sequence of (p, 3) float
          Intersection points, only returned if return_locations
        """

        # make sure input is _dtype for embree
        ray_origins = np.array(ray_origins, dtype=np.float64)
        ray_directions = np.array(ray_directions, dtype=np.float64)
        if ray_origins.shape != ray_directions.shape:
            raise ValueError("Ray origin and direction don't match!")
        ray_directions = util.unitize(ray_directions)

        # stack results for multiple hits into a sequence
        result_triangle = [np.zeros(0, dtype=np.int64)]
        result_ray_idx = [np.zeros(0, dtype=np.int64)]
        result_locations = [np.zeros((0, 3), dtype=np.float64)]

        if multiple_hits or return_locations:
            # how much to offset ray to transport to the other side of face
            ray_offsets = ray_directions * _ray_offset

            # grab the planes from triangles
            plane_origins = self.mesh.triangles[:, 0, :]
            plane_normals = self.mesh.face_normals

        # what each ray hit last iteration; -1 means nothing, used to
        # detect a ray stuck on the same face due to precision issues
        last_hit = np.full(len(ray_origins), -1, dtype=np.int64)
        # absolute indices of rays still being queried; shrinks each
        # iteration as rays miss, escape, or get culled
        live = np.arange(len(ray_origins))

        # use a for loop rather than a while to ensure this exits
        # if a ray is offset from a triangle and then is reported
        # hitting itself this could get stuck on that one triangle
        for _depth in range(max_hits):
            # if you set output=1 embreex returns distance along the ray
            # which is bizarrely slower than our own plane-line calc
            # TODO: switch `run(..., output=True)` once embreex>=4.4.0rc1 is stable
            query = self._scene.run(ray_origins[live], ray_directions[live])
            hit = query != -1
            if not hit.any():
                break

            # absolute indices and triangle indices for rays that hit
            hit_rays = live[hit]
            hit_tris = query[hit]

            # first-hit-only fast path: no duplicates or locations to track
            if not multiple_hits and not return_locations:
                result_triangle.append(hit_tris)
                result_ray_idx.append(hit_rays)
                break

            # rays that hit the same triangle as last iteration are stuck
            dupe = last_hit[hit_rays] == hit_tris
            # store the last hit so we can track duplicates
            last_hit[hit_rays] = hit_tris
            # subset to separate duplicates and non-duplicates
            ok_rays = hit_rays[~dupe]
            ok_tris = hit_tris[~dupe]
            dupe_rays = hit_rays[dupe]

            # compute where clean hits actually land on their triangle;
            # planes_lines silently drops near-parallel rays, so `valid`
            # shortens new_origins — we must trim ok_rays / ok_tris to match
            new_origins, valid = intersections.planes_lines(
                plane_origins=plane_origins[ok_tris],
                plane_normals=plane_normals[ok_tris],
                line_origins=ray_origins[ok_rays],
                line_directions=ray_directions[ok_rays],
            )
            ok_rays, ok_tris = ok_rays[valid], ok_tris[valid]

            result_locations.append(new_origins)
            result_triangle.append(ok_tris)
            result_ray_idx.append(ok_rays)

            if not multiple_hits:
                break

            # clean hits step onto the new face with a fresh base offset
            ray_origins[ok_rays] = new_origins + ray_offsets[ok_rays]
            ray_offsets[ok_rays] = ray_directions[ok_rays] * _ray_offset
            # stuck rays double their offset and step further to try to clear
            ray_offsets[dupe_rays] *= 2.0
            ray_origins[dupe_rays] += ray_offsets[dupe_rays]

            # carry forward only rays we successfully advanced;
            # dropped rays (misses, near-parallel planes) die here
            live = np.concatenate([ok_rays, dupe_rays])

        index_tri = np.concatenate(result_triangle)
        index_ray = np.concatenate(result_ray_idx)
        if return_locations:
            return index_tri, index_ray, np.concatenate(result_locations)
        return index_tri, index_ray

    @log_time
    def intersects_first(self, ray_origins, ray_directions):
        """
        Find the index of the first triangle a ray hits.


        Parameters
        ----------
        ray_origins : (n, 3) float
          Origins of rays
        ray_directions : (n, 3) float
          Direction (vector) of rays

        Returns
        ----------
        triangle_index : (n,) int
          Index of triangle ray hit, or -1 if not hit
        """
        # make sure our arrays are in the `embree` dtype
        ray_origins = np.array(ray_origins, dtype=_embree_dtype)
        ray_directions = np.array(ray_directions, dtype=_embree_dtype)
        if ray_origins.shape != ray_directions.shape:
            raise ValueError("Ray origin and direction don't match!")
        ray_directions = util.unitize(ray_directions)

        return self._scene.run(ray_origins, ray_directions)

    def intersects_any(self, ray_origins, ray_directions):
        """
        Check if a list of rays hits the surface.


        Parameters
        -----------
        ray_origins : (n, 3) float
          Origins of rays
        ray_directions : (n, 3) float
          Direction (vector) of rays

        Returns
        ----------
        hit : (n,) bool
          Did each ray hit the surface
        """

        first = self.intersects_first(
            ray_origins=ray_origins, ray_directions=ray_directions
        )
        hit = first != -1
        return hit

    def contains_points(self, points):
        """
        Check if a mesh contains a list of points, using ray tests.

        If the point is on the surface of the mesh, behavior is undefined.

        Parameters
        ---------
        points: (n, 3) points in space

        Returns
        ---------
        contains: (n,) bool
                         Whether point is inside mesh or not
        """
        return contains_points(self, points)

    def __getstate__(self):
        state = self.__dict__.copy()
        # don't pickle cache
        state.pop("_cache", None)
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        # Add cache back since it doesn't exist in the pickle
        self._cache = caching.Cache(id_function=self.mesh.__hash__)

    def __deepcopy__(self, *args):
        return self.__copy__()

    def __copy__(self, *args):
        return RayMeshIntersector(geometry=self.mesh, scale_to_box=self._scale_to_box)


class _EmbreeWrap:
    """
    A light wrapper for Embreex scene objects which
    allows queries to be scaled to help with precision
    issues, as well as selecting the correct dtypes.
    """

    def __init__(self, vertices, faces, scale):
        scaled = np.array(vertices, dtype=np.float64)
        self.origin = scaled.min(axis=0)
        self.scale = float(scale)
        scaled = (scaled - self.origin) * self.scale

        self.scene = rtcore_scene.EmbreeScene()
        # assign the geometry to the scene
        TriangleMesh(
            scene=self.scene,
            vertices=scaled.astype(_embree_dtype),
            indices=faces.view(np.ndarray).astype(np.int32),
        )

    def run(self, origins, normals, **kwargs):
        scaled = (np.array(origins, dtype=np.float64) - self.origin) * self.scale

        return self.scene.run(
            scaled.astype(_embree_dtype), normals.astype(_embree_dtype), **kwargs
        )
