# Copyright (c) 2024 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.


import numpy as np

import warp as wp
import trimesh

wp.set_module_options({"enable_backward": False})

NUM_QUERY_POINTS = 1000000

@wp.kernel
def sample_mesh_query(
    mesh: wp.uint64,
    query_points: wp.array(dtype=wp.vec3),
    query_d_max: float,
    query_closest_points: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    p = query_points[tid]
    query = wp.mesh_query_point_no_sign(mesh, p, query_d_max)

    if query.result:
        face = query.face
        cp = wp.vec3(float(face), query.u, query.v)
        query_closest_points[tid] = cp



if __name__ == '__main__':
    # mesh_path:
    mesh_root_path = ""
    mesh_path_crash = mesh_root_path +  "frame.stl"
    mesh_path_works = mesh_root_path + "active_rail.stl"


    mesh_path = mesh_path_crash
    device = wp.get_device("cuda:0")
    wp.config.mode = "debug"
    wp.config.verify_cuda = True
    wp.init()
    wp.build.clear_kernel_cache()
    wp.load_module(device=device)
    mesh = trimesh.load(mesh_path)

    points = np.array(mesh.vertices.view(np.ndarray))
    indices = mesh.faces.reshape(-1)
    bounding_box = np.array([points.min(axis=0), points.max(axis=0)])
    seed = 42

    rng = np.random.default_rng(seed)
    seed = seed + 1

    query_points_np = rng.uniform(bounding_box[0, :], bounding_box[1, :], size=(NUM_QUERY_POINTS, 3)).astype(
        np.float32
    )
    query_points = wp.array(query_points_np, dtype=wp.vec3, device=device)

    # create wp mesh
    mesh = wp.Mesh(
        points=wp.array(points, dtype=wp.vec3, device=device),
        velocities=None,
        indices=wp.array(indices, dtype=int, device=device),
    )

    query_closest_points = wp.empty_like(query_points, device=device)




    wp.synchronize_device(device)


    cmd = wp.launch(
            sample_mesh_query,
            dim=(NUM_QUERY_POINTS,),
            inputs=[mesh.id, query_points, 1000.0, query_closest_points],
            device=device,
            record_cmd=False,
        )


    wp.synchronize_device(device)