# Source - https://stackoverflow.com/q/79817960
# Posted by dw218192
# Retrieved 2025-11-13, License - CC BY-SA 4.0

import matplotlib.pyplot as plt
import numpy as np
import warp as wp


@wp.func
def get_triangle_points(
    mesh_id: wp.uint64, face_idx: int
) -> tuple[wp.vec3, wp.vec3, wp.vec3]:
    i1 = wp.mesh_get_index(mesh_id, 3 * face_idx)
    i2 = wp.mesh_get_index(mesh_id, 3 * face_idx + 1)
    i3 = wp.mesh_get_index(mesh_id, 3 * face_idx + 2)
    v1 = wp.mesh_get_point(mesh_id, i1)
    v2 = wp.mesh_get_point(mesh_id, i2)
    v3 = wp.mesh_get_point(mesh_id, i3)
    return v1, v2, v3


@wp.kernel
def render_quad(
    mesh_id: wp.uint64,
    width: int,
    height: int,
    pixels: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    x = tid % width
    y = tid // width

    sx = 2.0 * (float(x) + 0.5) / float(width) - 1.0
    sy = 2.0 * (float(y) + 0.5) / float(height) - 1.0

    ro = wp.vec3(0.0, 5.0, 0.0)
    rd = wp.normalize(wp.vec3(-1.0, sy, sx))

    query = wp.mesh_query_ray(mesh_id, ro, rd, 1.0e6)

    color = wp.vec3(0.0, 0.0, 0.0)

    if query.result:
        face_idx = query.face
        v0, v1, v2 = get_triangle_points(mesh_id, face_idx)
        area = 0.5 * wp.length(wp.cross(v1 - v0, v2 - v0))
        if wp.abs(area) < 1e-5:
            color = wp.vec3(0.0, 1.0, 0.0)
        else:
            color = wp.vec3(1.0, 0.0, 0.0)

    pixels[tid] = color


def main():
    quad_points = np.array(
        [
            [-5.0, 0.0, 5.0],
            [-5.0, 0.0, -5.0],
            [-5.0, 10.0, -5.0],
            [-5.0, 10.0, 5.0],
        ],
        dtype=np.float32,
    )

    quad_indices = np.array([0, 1, 2, 0, 2, 3], dtype=np.int32)
    wp_points = wp.array(np.ascontiguousarray(quad_points), dtype=wp.vec3)
    wp_indices = wp.array(np.ascontiguousarray(quad_indices), dtype=wp.int32)
    wp_mesh = wp.Mesh(points=wp_points, indices=wp_indices, velocities=None)

    width = 256
    height = 256

    pixels = wp.zeros(width * height, dtype=wp.vec3)

    wp.launch(
        kernel=render_quad,
        dim=width * height,
        inputs=[wp_mesh.id, width, height, pixels],
    )

    image = pixels.numpy().reshape((height, width, 3))

    plt.figure(figsize=(6, 6))
    plt.imshow(image, origin="lower", interpolation="nearest")
    plt.axis("off")
    plt.show()


if __name__ == "__main__":
    main()