# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "jax[cuda]>=0.7.0",
#     "jaxtyping",
# ]
# ///

import jax
import jax.numpy as jnp
from jaxtyping import Array, Bool, Float


def ray_intersect_triangle(
    ray_origin: Float[Array, "3"],
    ray_direction: Float[Array, "3"],
    triangle_vertices: Float[Array, "3 3"],
) -> Bool[Array, " "]:
    # Placeholder computations for the Möller-Trumbore algorithm
    t = (ray_origins * ray_direction).sum()
    hit = (ray_origin * ray_direction * triangle_vertices[0, :]).sum() > 0
    # Actual ray-intersection boolean value
    return (t <= 1.0) & hit


# Intermediate function, used to show that we perform a reduction on 'num_triangles'
@jax.jit
def _ray_intersect_any_triangle(
    ray_origin: Float[Array, "3"],
    ray_direction: Float[Array, "3"],
    triangle_vertices: Float[Array, "num_triangles 3 3"],
) -> Bool[Array, " "]:
    return jax.vmap(ray_intersect_triangle, in_axes=(None, None, 0))(
        ray_origin, ray_direction, triangle_vertices
    ).any()


# Vectorized function, that is the bottleneck
@jax.jit
def rays_intersect_any_triangle(
    ray_origins: Float[Array, "num_rays 3"],
    ray_directions: Float[Array, "num_rays 3"],
    triangle_vertices: Float[Array, "num_triangles 3 3"],
) -> Bool[Array, "num_rays"]:
    # Solution 1: OOM
    return jnp.vectorize(
        _ray_intersect_any_triangle, excluded={2}, signature="(3),(3)->()"
    )(ray_origins, ray_directions, triangle_vertices)
    # Solution 2: OOM
    return jax.vmap(_ray_intersect_any_triangle, in_axes=(0, 0, None))(
        ray_origins, ray_directions, triangle_vertices
    )
    # Solution 3: takes ages to compute
    return jax.lax.map(
        lambda xs: _ray_intersect_any_triangle(*xs, triangle_vertices),
        (ray_origins, ray_directions),
    )

    # Solution 4: OOM (but uses way less memory), and super slow if we reduce 'num_rays'
    def scan_fun(
        intersect: Bool[Array, "num_rays"], triangle_vertices: Float[Array, "3 3"]
    ) -> tuple[Bool[Array, "num_rays"], None]:
        intersect |= jax.vmap(ray_intersect_triangle, in_axes=(0, 0, None))(
            ray_origins, ray_directions, triangle_vertices
        )
        return intersect, None

    return jax.lax.scan(
        scan_fun,
        init=jnp.zeros(ray_origins.shape[0], dtype=bool),
        xs=triangle_vertices,
    )


if __name__ == "__main__":
    key = jax.random.key(1234)

    num_rays = 1_000_000
    num_triangles = 1_000_000

    key_ray_origins, key_ray_directions, key_triangle_vertices = jax.random.split(
        key, 3
    )
    ray_origins = jax.random.normal(key_ray_origins, (num_rays, 3))
    ray_directions = jax.random.normal(key_ray_directions, (num_rays, 3))
    triangle_vertices = jax.random.normal(key_triangle_vertices, (num_triangles, 3, 3))

    # Boom
    rays_intersect_any_triangle(ray_origins, ray_directions, triangle_vertices)