# 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.

import warp as wp
import numpy as np
import math


def _usd_add_xform(prim):
    from pxr import UsdGeom

    prim = UsdGeom.Xform(prim)
    prim.ClearXformOpOrder()

    t = prim.AddTranslateOp()
    r = prim.AddOrientOp()
    s = prim.AddScaleOp()


def _usd_set_xform(xform, pos: tuple, rot: tuple, scale: tuple, time):
    from pxr import UsdGeom, Gf

    xform = UsdGeom.Xform(xform)

    xform_ops = xform.GetOrderedXformOps()

    xform_ops[0].Set(Gf.Vec3d(float(pos[0]), float(pos[1]), float(pos[2])), time)
    xform_ops[1].Set(Gf.Quatf(float(rot[3]), float(rot[0]), float(rot[1]), float(rot[2])), time)
    xform_ops[2].Set(Gf.Vec3d(float(scale[0]), float(scale[1]), float(scale[2])), time)


# transforms a cylinder such that it connects the two points pos0, pos1
def _compute_segment_xform(pos0, pos1):
    from pxr import Gf

    mid = (pos0 + pos1) * 0.5
    height = (pos1 - pos0).GetLength()

    dir = (pos1 - pos0) / height

    rot = Gf.Rotation()
    rot.SetRotateInto((0.0, 0.0, 1.0), Gf.Vec3d(dir))

    scale = Gf.Vec3f(1.0, 1.0, height)

    return (mid, Gf.Quath(rot.GetQuat()), scale)


class UsdRenderer:
    """A USD renderer"""

    def __init__(self, stage, up_axis="Y", fps=60, scaling=1.0):
        """Construct a UsdRenderer object

        Args:
            model: A simulation model
            stage (str/Usd.Stage): A USD stage (either in memory or on disk)
            up_axis (str): The upfacing axis of the stage
            fps: The number of frames per second to use in the USD file
            scaling: Scaling factor to use for the entities in the scene
        """

        from pxr import Usd, UsdGeom, UsdLux, Sdf, Gf

        if isinstance(stage, str):
            self.stage = stage = Usd.Stage.CreateNew(stage)
        elif isinstance(stage, Usd.Stage):
            self.stage = stage
        else:
            print("Failed to create stage in renderer. Please construct with stage path or stage object.")
        self.up_axis = up_axis.upper()
        self.fps = float(fps)
        self.time = 0.0

        self.draw_points = True
        self.draw_springs = False
        self.draw_triangles = False

        self.root = UsdGeom.Xform.Define(stage, "/root")

        # mapping from shape ID to UsdGeom class
        self._shape_constructors = {}
        # optional scaling applied to shape instances (e.g. cubes)
        self._shape_custom_scale = {}

        # apply scaling
        self.root.ClearXformOpOrder()
        s = self.root.AddScaleOp()
        s.Set(Gf.Vec3d(float(scaling), float(scaling), float(scaling)), 0.0)

        self.stage.SetDefaultPrim(self.root.GetPrim())
        self.stage.SetStartTimeCode(0.0)
        self.stage.SetEndTimeCode(0.0)
        self.stage.SetTimeCodesPerSecond(self.fps)

        if up_axis == "X":
            UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.x)
        elif up_axis == "Y":
            UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.y)
        elif up_axis == "Z":
            UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)

        # add default lights
        light_0 = UsdLux.DistantLight.Define(stage, "/light_0")
        light_0.GetPrim().CreateAttribute("intensity", Sdf.ValueTypeNames.Float, custom=False).Set(2500.0)
        light_0.GetPrim().CreateAttribute("color", Sdf.ValueTypeNames.Color3f, custom=False).Set(
            Gf.Vec3f(0.98, 0.85, 0.7)
        )

        UsdGeom.Xform(light_0.GetPrim()).AddRotateYOp().Set(value=(70.0))
        UsdGeom.Xform(light_0.GetPrim()).AddRotateXOp().Set(value=(-45.0))

        light_1 = UsdLux.DistantLight.Define(stage, "/light_1")
        light_1.GetPrim().CreateAttribute("intensity", Sdf.ValueTypeNames.Float, custom=False).Set(2500.0)
        light_1.GetPrim().CreateAttribute("color", Sdf.ValueTypeNames.Color3f, custom=False).Set(
            Gf.Vec3f(0.62, 0.82, 0.98)
        )

        UsdGeom.Xform(light_1.GetPrim()).AddRotateYOp().Set(value=(-70.0))
        UsdGeom.Xform(light_1.GetPrim()).AddRotateXOp().Set(value=(-45.0))

    def begin_frame(self, time):
        self.stage.SetEndTimeCode(time * self.fps)
        self.time = time * self.fps

    def end_frame(self):
        pass

    def register_body(self, body_name):
        from pxr import UsdGeom

        xform = UsdGeom.Xform.Define(self.stage, self.root.GetPath().AppendChild(body_name))

        _usd_add_xform(xform)

    def _resolve_path(self, name, parent_body=None, is_template=False):
        # resolve the path to the prim with the given name and optional parent body
        if is_template:
            return self.root.GetPath().AppendChild("_template_shapes").AppendChild(name)
        if parent_body is None:
            return self.root.GetPath().AppendChild(name)
        else:
            return self.root.GetPath().AppendChild(parent_body).AppendChild(name)

    def add_shape_instance(
        self,
        name: str,
        shape,
        body,
        pos: tuple,
        rot: tuple,
        scale: tuple = (1.0, 1.0, 1.0),
        color: tuple = (1.0, 1.0, 1.0),
    ):
        sdf_path = self._resolve_path(name, body)
        instance = self._shape_constructors[shape.name].Define(self.stage, sdf_path)
        instance.GetPrim().GetReferences().AddInternalReference(shape)

        _usd_add_xform(instance)
        if shape.name in self._shape_custom_scale:
            cs = self._shape_custom_scale[shape.name]
            scale = (scale[0] * cs[0], scale[1] * cs[1], scale[2] * cs[2])
        _usd_set_xform(instance, pos, rot, scale, self.time)

    def render_plane(
        self,
        name: str,
        pos: tuple,
        rot: tuple,
        width: float,
        length: float,
        color: tuple = None,
        parent_body: str = None,
        is_template: bool = False,
    ):
        """
        Render a plane with the given dimensions.

        Args:
            name: Name of the plane
            pos: Position of the plane
            rot: Rotation of the plane
            width: Width of the plane
            length: Length of the plane
            color: Color of the plane
            parent_body: Name of the parent body
            is_template: Whether the plane is a template
        """
        from pxr import UsdGeom, Sdf

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            plane_path = prim_path.AppendChild("plane")
        else:
            plane_path = self._resolve_path(name, parent_body)
            prim_path = plane_path

        plane = UsdGeom.Mesh.Get(self.stage, plane_path)
        if not plane:
            plane = UsdGeom.Mesh.Define(self.stage, plane_path)
            plane.CreateDoubleSidedAttr().Set(True)

            width = width if width > 0.0 else 100.0
            length = length if length > 0.0 else 100.0
            points = ((-width, 0.0, -length), (width, 0.0, -length), (width, 0.0, length), (-width, 0.0, length))
            normals = ((0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0))
            counts = (4,)
            indices = [0, 1, 2, 3]

            plane.GetPointsAttr().Set(points)
            plane.GetNormalsAttr().Set(normals)
            plane.GetFaceVertexCountsAttr().Set(counts)
            plane.GetFaceVertexIndicesAttr().Set(indices)
            _usd_add_xform(plane)

        self._shape_constructors[name] = UsdGeom.Mesh

        if not is_template:
            _usd_set_xform(plane, pos, rot, (1.0, 1.0, 1.0), 0.0)

        return prim_path

    def render_ground(self, size: float = 100.0):
        from pxr import UsdGeom

        mesh = UsdGeom.Mesh.Define(self.stage, self.root.GetPath().AppendChild("ground"))
        mesh.CreateDoubleSidedAttr().Set(True)

        if self.up_axis == "X":
            points = ((0.0, -size, -size), (0.0, size, -size), (0.0, size, size), (0.0, -size, size))
            normals = ((1.0, 0.0, 0.0), (1.0, 0.0, 0.0), (1.0, 0.0, 0.0), (1.0, 0.0, 0.0))
        elif self.up_axis == "Y":
            points = ((-size, 0.0, -size), (size, 0.0, -size), (size, 0.0, size), (-size, 0.0, size))
            normals = ((0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0))
        elif self.up_axis == "Z":
            points = ((-size, -size, 0.0), (size, -size, 0.0), (size, size, 0.0), (-size, size, 0.0))
            normals = ((0.0, 0.0, 1.0), (0.0, 0.0, 1.0), (0.0, 0.0, 1.0), (0.0, 0.0, 1.0))
        counts = (4,)
        indices = [0, 1, 2, 3]

        mesh.GetPointsAttr().Set(points)
        mesh.GetNormalsAttr().Set(normals)
        mesh.GetFaceVertexCountsAttr().Set(counts)
        mesh.GetFaceVertexIndicesAttr().Set(indices)

    def render_sphere(
        self, name: str, pos: tuple, rot: tuple, radius: float, parent_body: str = None, is_template: bool = False
    ):
        """Debug helper to add a sphere for visualization

        Args:
            pos: The position of the sphere
            radius: The radius of the sphere
            name: A name for the USD prim on the stage
        """

        from pxr import UsdGeom, Sdf

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            sphere_path = prim_path.AppendChild("sphere")
        else:
            sphere_path = self._resolve_path(name, parent_body)
            prim_path = sphere_path

        sphere = UsdGeom.Sphere.Get(self.stage, sphere_path)
        if not sphere:
            sphere = UsdGeom.Sphere.Define(self.stage, sphere_path)
            _usd_add_xform(sphere)

        sphere.GetRadiusAttr().Set(radius, self.time)

        self._shape_constructors[name] = UsdGeom.Sphere

        if not is_template:
            _usd_set_xform(sphere, pos, rot, (1.0, 1.0, 1.0), 0.0)

        return prim_path

    def render_capsule(
        self,
        name: str,
        pos: tuple,
        rot: tuple,
        radius: float,
        half_height: float,
        parent_body: str = None,
        is_template: bool = False,
    ):
        """
        Debug helper to add a capsule for visualization

        Args:
            pos: The position of the capsule
            radius: The radius of the capsule
            half_height: The half height of the capsule
            name: A name for the USD prim on the stage
        """

        from pxr import UsdGeom, Sdf

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            capsule_path = prim_path.AppendChild("capsule")
        else:
            capsule_path = self._resolve_path(name, parent_body)
            prim_path = capsule_path

        capsule = UsdGeom.Capsule.Get(self.stage, capsule_path)
        if not capsule:
            capsule = UsdGeom.Capsule.Define(self.stage, capsule_path)
            _usd_add_xform(capsule)

        capsule.GetRadiusAttr().Set(float(radius))
        capsule.GetHeightAttr().Set(float(half_height * 2.0))
        capsule.GetAxisAttr().Set("Y")

        self._shape_constructors[name] = UsdGeom.Capsule

        if not is_template:
            _usd_set_xform(capsule, pos, rot, (1.0, 1.0, 1.0), 0.0)

        return prim_path

    def render_cylinder(
        self,
        name: str,
        pos: tuple,
        rot: tuple,
        radius: float,
        half_height: float,
        parent_body: str = None,
        is_template: bool = False,
    ):
        """
        Debug helper to add a cylinder for visualization

        Args:
            pos: The position of the cylinder
            radius: The radius of the cylinder
            half_height: The half height of the cylinder
            name: A name for the USD prim on the stage
        """

        from pxr import UsdGeom, Sdf

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            cylinder_path = prim_path.AppendChild("cylinder")
        else:
            cylinder_path = self._resolve_path(name, parent_body)
            prim_path = cylinder_path

        cylinder = UsdGeom.Cylinder.Get(self.stage, cylinder_path)
        if not cylinder:
            cylinder = UsdGeom.Cylinder.Define(self.stage, cylinder_path)
            _usd_add_xform(cylinder)

        cylinder.GetRadiusAttr().Set(float(radius))
        cylinder.GetHeightAttr().Set(float(half_height * 2.0))
        cylinder.GetAxisAttr().Set("Y")

        self._shape_constructors[name] = UsdGeom.Cylinder

        if not is_template:
            _usd_set_xform(cylinder, pos, rot, (1.0, 1.0, 1.0), 0.0)

        return prim_path

    def render_cone(
        self,
        name: str,
        pos: tuple,
        rot: tuple,
        radius: float,
        half_height: float,
        parent_body: str = None,
        is_template: bool = False,
    ):
        """
        Debug helper to add a cone for visualization

        Args:
            pos: The position of the cone
            radius: The radius of the cone
            half_height: The half height of the cone
            name: A name for the USD prim on the stage
        """

        from pxr import UsdGeom, Sdf

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            cone_path = prim_path.AppendChild("cone")
        else:
            cone_path = self._resolve_path(name, parent_body)
            prim_path = cone_path

        cone = UsdGeom.Cone.Get(self.stage, cone_path)
        if not cone:
            cone = UsdGeom.Cone.Define(self.stage, cone_path)
            _usd_add_xform(cone)

        cone.GetRadiusAttr().Set(float(radius))
        cone.GetHeightAttr().Set(float(half_height * 2.0))
        cone.GetAxisAttr().Set("Y")

        self._shape_constructors[name] = UsdGeom.Cone

        if not is_template:
            _usd_set_xform(cone, pos, rot, (1.0, 1.0, 1.0), 0.0)

        return prim_path

    def render_box(
        self, name: str, pos: tuple, rot: tuple, extents: tuple, parent_body: str = None, is_template: bool = False
    ):
        """Debug helper to add a box for visualization

        Args:
            pos: The position of the sphere
            extents: The radius of the sphere
            name: A name for the USD prim on the stage
        """

        from pxr import UsdGeom, Sdf, Gf, Vt

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            cube_path = prim_path.AppendChild("cube")
        else:
            cube_path = self._resolve_path(name, parent_body)
            prim_path = cube_path

        cube = UsdGeom.Cube.Get(self.stage, cube_path)
        if not cube:
            cube = UsdGeom.Cube.Define(self.stage, cube_path)
            _usd_add_xform(cube)

        self._shape_constructors[name] = UsdGeom.Cube
        self._shape_custom_scale[name] = extents

        if not is_template:
            _usd_set_xform(cube, pos, rot, extents, 0.0)

        return prim_path

    def render_ref(self, name: str, path: str, pos: tuple, rot: tuple, scale: tuple):
        from pxr import UsdGeom

        ref_path = "/root/" + name

        ref = UsdGeom.Xform.Get(self.stage, ref_path)
        if not ref:
            ref = UsdGeom.Xform.Define(self.stage, ref_path)
            ref.GetPrim().GetReferences().AddReference(path)
            _usd_add_xform(ref)

        # update transform
        _usd_set_xform(ref, pos, rot, scale, self.time)

    def render_mesh(
        self,
        name: str,
        points,
        indices,
        colors=None,
        pos=(0.0, 0.0, 0.0),
        rot=(0.0, 0.0, 0.0, 1.0),
        scale=(1.0, 1.0, 1.0),
        update_topology=False,
        parent_body: str = None,
        is_template: bool = False,
    ):
        from pxr import UsdGeom, Sdf

        if is_template:
            prim_path = self._resolve_path(name, parent_body, is_template)
            blueprint = UsdGeom.Scope.Define(self.stage, prim_path)
            blueprint_prim = blueprint.GetPrim()
            blueprint_prim.SetInstanceable(True)
            blueprint_prim.SetSpecifier(Sdf.SpecifierClass)
            mesh_path = prim_path.AppendChild("mesh")
        else:
            mesh_path = self._resolve_path(name, parent_body)
            prim_path = mesh_path

        mesh = UsdGeom.Mesh.Get(self.stage, mesh_path)
        if not mesh:
            mesh = UsdGeom.Mesh.Define(self.stage, mesh_path)
            UsdGeom.Primvar(mesh.GetDisplayColorAttr()).SetInterpolation("vertex")
            _usd_add_xform(mesh)

            # force topology update on first frame
            update_topology = True

        mesh.GetPointsAttr().Set(points, self.time)

        if update_topology:
            idxs = np.array(indices).reshape(-1, 3)
            mesh.GetFaceVertexIndicesAttr().Set(idxs, self.time)
            mesh.GetFaceVertexCountsAttr().Set([3] * len(idxs), self.time)

        if colors:
            mesh.GetDisplayColorAttr().Set(colors, self.time)

        self._shape_constructors[name] = UsdGeom.Mesh
        self._shape_custom_scale[name] = scale

        if not is_template:
            _usd_set_xform(mesh, pos, rot, scale, self.time)

        return prim_path

    def render_line_list(self, name, vertices, indices, color, radius):
        """Debug helper to add a line list as a set of capsules

        Args:
            vertices: The vertices of the line-strip
            color: The color of the line
            time: The time to update at
        """

        from pxr import UsdGeom, Gf

        num_lines = int(len(indices) / 2)

        if num_lines < 1:
            return

        # look up rope point instancer
        instancer_path = self.root.GetPath().AppendChild(name)
        instancer = UsdGeom.PointInstancer.Get(self.stage, instancer_path)

        if not instancer:
            instancer = UsdGeom.PointInstancer.Define(self.stage, instancer_path)
            instancer_capsule = UsdGeom.Capsule.Define(self.stage, instancer.GetPath().AppendChild("capsule"))
            instancer_capsule.GetRadiusAttr().Set(radius)
            instancer.CreatePrototypesRel().SetTargets([instancer_capsule.GetPath()])
            # instancer.CreatePrimvar("displayColor", Sdf.ValueTypeNames.Float3Array, "constant", 1)

        line_positions = []
        line_rotations = []
        line_scales = []

        for i in range(num_lines):
            pos0 = vertices[indices[i * 2 + 0]]
            pos1 = vertices[indices[i * 2 + 1]]

            (pos, rot, scale) = _compute_segment_xform(
                Gf.Vec3f(float(pos0[0]), float(pos0[1]), float(pos0[2])),
                Gf.Vec3f(float(pos1[0]), float(pos1[1]), float(pos1[2])),
            )

            line_positions.append(pos)
            line_rotations.append(rot)
            line_scales.append(scale)
            # line_colors.append(Gf.Vec3f((float(i)/num_lines, 0.5, 0.5)))

        instancer.GetPositionsAttr().Set(line_positions, self.time)
        instancer.GetOrientationsAttr().Set(line_rotations, self.time)
        instancer.GetScalesAttr().Set(line_scales, self.time)
        instancer.GetProtoIndicesAttr().Set([0] * num_lines, self.time)

    #      instancer.GetPrimvar("displayColor").Set(line_colors, time)

    def render_line_strip(self, name: str, vertices, color: tuple, radius: float = 0.01):
        from pxr import UsdGeom, Gf

        num_lines = int(len(vertices) - 1)

        if num_lines < 1:
            return

        # look up rope point instancer
        instancer_path = self.root.GetPath().AppendChild(name)
        instancer = UsdGeom.PointInstancer.Get(self.stage, instancer_path)

        if not instancer:
            instancer = UsdGeom.PointInstancer.Define(self.stage, instancer_path)
            instancer_capsule = UsdGeom.Capsule.Define(self.stage, instancer.GetPath().AppendChild("capsule"))
            instancer_capsule.GetRadiusAttr().Set(radius)
            instancer.CreatePrototypesRel().SetTargets([instancer_capsule.GetPath()])

        line_positions = []
        line_rotations = []
        line_scales = []

        for i in range(num_lines):
            pos0 = vertices[i]
            pos1 = vertices[i + 1]

            (pos, rot, scale) = _compute_segment_xform(
                Gf.Vec3f(float(pos0[0]), float(pos0[1]), float(pos0[2])),
                Gf.Vec3f(float(pos1[0]), float(pos1[1]), float(pos1[2])),
            )

            line_positions.append(pos)
            line_rotations.append(rot)
            line_scales.append(scale)

        instancer.GetPositionsAttr().Set(line_positions, self.time)
        instancer.GetOrientationsAttr().Set(line_rotations, self.time)
        instancer.GetScalesAttr().Set(line_scales, self.time)
        instancer.GetProtoIndicesAttr().Set([0] * num_lines, self.time)

        instancer_capsule = UsdGeom.Capsule.Get(self.stage, instancer.GetPath().AppendChild("capsule"))
        instancer_capsule.GetDisplayColorAttr().Set([Gf.Vec3f(color)], self.time)

    def render_points(self, name: str, points, radius, colors=None):
        from pxr import UsdGeom, Gf

        instancer_path = self.root.GetPath().AppendChild(name)
        instancer = UsdGeom.PointInstancer.Get(self.stage, instancer_path)
        radius_is_scalar = np.isscalar(radius)
        if not instancer:
            if colors is None:
                instancer = UsdGeom.PointInstancer.Define(self.stage, instancer_path)
                instancer_sphere = UsdGeom.Sphere.Define(self.stage, instancer.GetPath().AppendChild("sphere"))
                if radius_is_scalar:
                    instancer_sphere.GetRadiusAttr().Set(radius)
                else:
                    instancer_sphere.GetRadiusAttr().Set(1.0)
                    instancer.GetScalesAttr().Set(np.tile(radius, (3, 1)).T)

                instancer.CreatePrototypesRel().SetTargets([instancer_sphere.GetPath()])
                instancer.CreateProtoIndicesAttr().Set([0] * len(points))

                # set identity rotations
                quats = [Gf.Quath(1.0, 0.0, 0.0, 0.0)] * len(points)
                instancer.GetOrientationsAttr().Set(quats, self.time)
            else:
                from pxr import Sdf

                instancer = UsdGeom.Points.Define(self.stage, instancer_path)

                instancer.CreatePrimvar("displayColor", Sdf.ValueTypeNames.Float3Array, "vertex", 1)
                if radius_is_scalar:
                    instancer.GetWidthsAttr().Set([radius] * len(points))
                else:
                    instancer.GetWidthsAttr().Set(radius)

        if colors is None:
            instancer.GetPositionsAttr().Set(points, self.time)
        else:
            instancer.GetPointsAttr().Set(points, self.time)
            instancer.GetDisplayColorAttr().Set(colors, self.time)

    def update_body_transforms(self, body_q):
        from pxr import UsdGeom, Sdf

        if isinstance(body_q, wp.array):
            body_q = body_q.numpy()

        with Sdf.ChangeBlock():
            for b in range(self.model.body_count):
                node_name = self.body_names[b]
                node = UsdGeom.Xform(self.stage.GetPrimAtPath(self.root.GetPath().AppendChild(node_name)))

                # unpack rigid transform
                X_sb = wp.transform_expand(body_q[b])

                _usd_set_xform(node, X_sb.p, X_sb.q, (1.0, 1.0, 1.0), self.time)

    def save(self):
        try:
            self.stage.Save()
            return True
        except:
            print("Failed to save USD stage")
            return False