# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

#############################################################################
# Example Mesh Intersection
#
# Show how to use built-in BVH query to test if two triangle meshes intersect.
#
##############################################################################

import os

import numpy as np
from pxr import Usd, UsdGeom

import warp as wp
import warp.render

np.random.seed(42)


wp.init()


@wp.func
def cw_min(a: wp.vec3, b: wp.vec3):
    return wp.vec3(wp.min(a[0], b[0]), wp.min(a[1], b[1]), wp.min(a[2], b[2]))


@wp.func
def cw_max(a: wp.vec3, b: wp.vec3):
    return wp.vec3(wp.max(a[0], b[0]), wp.max(a[1], b[1]), wp.max(a[2], b[2]))


@wp.kernel
def intersect(
    mesh_0: wp.uint64,
    mesh_1: wp.uint64,
    num_faces: int,
    xforms: wp.array(dtype=wp.transform),
    result: wp.array(dtype=int),
):
    tid = wp.tid()

    # mesh_0 is assumed to be the query mesh, we launch one thread
    # for each face in mesh_0 and test it against the opposing mesh's BVH
    face = tid % num_faces
    batch = tid // num_faces

    # transforms from mesh_0 -> mesh_1 space
    xform = xforms[batch]

    # load query triangles points and transform to mesh_1's space
    v0 = wp.transform_point(xform, wp.mesh_eval_position(mesh_0, face, 1.0, 0.0))
    v1 = wp.transform_point(xform, wp.mesh_eval_position(mesh_0, face, 0.0, 1.0))
    v2 = wp.transform_point(xform, wp.mesh_eval_position(mesh_0, face, 0.0, 0.0))

    # compute bounds of the query triangle
    lower = cw_min(cw_min(v0, v1), v2)
    upper = cw_max(cw_max(v0, v1), v2)

    query = wp.mesh_query_aabb(mesh_1, lower, upper)

    for f in query:
        u0 = wp.mesh_eval_position(mesh_1, f, 1.0, 0.0)
        u1 = wp.mesh_eval_position(mesh_1, f, 0.0, 1.0)
        u2 = wp.mesh_eval_position(mesh_1, f, 0.0, 0.0)

        # test for triangle intersection
        i = wp.intersect_tri_tri(v0, v1, v2, u0, u1, u2)

        if i > 0:
            result[batch] = 1
            return

        # use if you want to count all intersections
        # wp.atomic_add(result, batch, i)


class Example:
    def __init__(self, stage):
        rng = np.random.default_rng()

        self.device = wp.get_device()
        self.query_count = 1024
        self.has_queried = False

        self.renderer = wp.render.UsdRenderer(stage)

        self.path_0 = "assets/cube.usda"
        self.path_1 = "assets/sphere.usda"

        self.mesh_0 = self.load_mesh(self.path_0, "/Cube/Cube_001")
        self.mesh_1 = self.load_mesh(self.path_1, "/Sphere/Sphere")

        self.query_num_faces = int(len(self.mesh_0.indices) / 3)
        self.query_num_points = len(self.mesh_0.points)

        # generate random relative transforms
        self.xforms = []

        for i in range(self.query_count):
            # random offset
            p = wp.vec3(rng.random(3) * 0.5 - 0.5) * 5.0

            # random orientation
            axis = wp.normalize(wp.vec3(rng.random(3) * 0.5 - 0.5))
            angle = float(np.random.rand(1)[0])

            q = wp.quat_from_axis_angle(wp.normalize(axis), angle)

            self.xforms.append(wp.transform(p, q))

        self.array_result = wp.zeros(self.query_count, dtype=int)
        self.array_xforms = wp.array(self.xforms, dtype=wp.transform)

        # compile and load the module up front (for accurate profiling)
        wp.load_module(device=self.device)

    def update(self):
        with wp.ScopedTimer("intersect", active=True):
            wp.launch(
                kernel=intersect,
                dim=self.query_num_faces * self.query_count,
                inputs=[self.mesh_0.id, self.mesh_1.id, self.query_num_faces, self.array_xforms, self.array_result],
            )

    def render(self):
        # bring results back to host
        result = self.array_result.numpy()

        with wp.ScopedTimer("render", active=True):
            self.renderer.begin_frame(0.0)

            for i in range(self.query_count):
                spacing = 8.0
                offset = i * spacing

                xform = self.xforms[i]
                self.renderer.render_ref(
                    f"mesh_{i}_0",
                    os.path.join(os.path.dirname(__file__), self.path_0),
                    pos=wp.vec3(xform.p[0] + offset, xform.p[1], xform.p[2]),
                    rot=xform.q,
                    scale=wp.vec3(1.0, 1.0, 1.0),
                )
                self.renderer.render_ref(
                    f"mesh_{i}_1",
                    os.path.join(os.path.dirname(__file__), self.path_1),
                    pos=wp.vec3(offset, 0.0, 0.0),
                    rot=wp.quat_identity(),
                    scale=wp.vec3(1.0, 1.0, 1.0),
                )

                # if pair intersects then draw a small box above the pair
                if result[i] > 0:
                    self.renderer.render_box(
                        f"result_{i}",
                        pos=wp.vec3(xform.p[0] + offset, xform.p[1] + 5.0, xform.p[2]),
                        rot=wp.quat_identity(),
                        extents=(0.1, 0.1, 0.1),
                    )

            self.renderer.end_frame()

    # create collision meshes
    def load_mesh(self, path, prim):
        usd_path = os.path.join(os.path.dirname(__file__), path)
        usd_stage = Usd.Stage.Open(usd_path)
        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath(prim))

        mesh = wp.Mesh(
            points=wp.array(usd_geom.GetPointsAttr().Get(), dtype=wp.vec3),
            indices=wp.array(usd_geom.GetFaceVertexIndicesAttr().Get(), dtype=int),
        )

        return mesh


if __name__ == "__main__":
    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_mesh_intersect.usd")

    example = Example(stage_path)

    example.update()
    example.render()

    example.renderer.save()