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

# Standard Library
import math
from typing import Dict, List, Optional, Union

# Third Party
import numpy as np
import torch
from tqdm import tqdm

# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
from curobo.geom.types import (
    Capsule,
    Cuboid,
    Cylinder,
    Material,
    Mesh,
    Obstacle,
    Sphere,
    WorldConfig,
)
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
from curobo.util.logger import log_error, log_info, log_warn
from curobo.util_file import (
    file_exists,
    get_assets_path,
    get_filename,
    get_files_from_dir,
    get_robot_configs_path,
    join_path,
    load_yaml,
)
from curobo.wrap.reacher.motion_gen import MotionGenResult

try:
    # Third Party
    from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics, UsdShade
except ImportError:
    raise ImportError(
        "usd-core failed to import, install with pip install usd-core"
        + " NOTE: Do not install this if using with ISAAC SIM."
    )


def set_prim_translate(prim, translation):
    UsdGeom.Xformable(prim).AddTranslateOp().Set(Gf.Vec3d(translation))


def set_prim_transform(
    prim, pose: List[float], scale: List[float] = [1, 1, 1], use_float: bool = False
):
    if not prim.GetAttribute("xformOp:translate").IsValid():
        UsdGeom.Xformable(prim).AddTranslateOp(UsdGeom.XformOp.PrecisionFloat)

    if prim.GetAttribute("xformOp:orient").IsValid():
        if isinstance(prim.GetAttribute("xformOp:orient").Get(), Gf.Quatf):
            use_float = True
    else:
        UsdGeom.Xformable(prim).AddOrientOp(UsdGeom.XformOp.PrecisionFloat)
        use_float = True

    if not prim.GetAttribute("xformOp:scale").IsValid():
        UsdGeom.Xformable(prim).AddScaleOp(UsdGeom.XformOp.PrecisionFloat)
    quat = pose[3:]

    if use_float:
        position = Gf.Vec3f(pose[:3])
        q = Gf.Quatf(quat[0], quat[1:])
        dims = Gf.Vec3f(scale)

    else:
        position = Gf.Vec3d(pose[:3])
        q = Gf.Quatd(quat[0], quat[1:])
        dims = Gf.Vec3d(scale)

    prim.GetAttribute("xformOp:translate").Set(position)
    prim.GetAttribute("xformOp:orient").Set(q)
    prim.GetAttribute("xformOp:scale").Set(dims)

    # get scale:


def get_prim_world_pose(cache: UsdGeom.XformCache, prim: Usd.Prim, inverse: bool = False):
    world_transform: Gf.Matrix4d = cache.GetLocalToWorldTransform(prim)
    # get scale:
    scale: Gf.Vec3d = Gf.Vec3d(*(v.GetLength() for v in world_transform.ExtractRotationMatrix()))
    scale = list(scale)
    t_mat = world_transform.RemoveScaleShear()
    if inverse:
        t_mat = t_mat.GetInverse()

    # mat = np.zeros((4,4))
    # mat[:,:] = t_mat
    translation: Gf.Vec3d = t_mat.ExtractTranslation()
    rotation: Gf.Rotation = t_mat.ExtractRotation()
    q = rotation.GetQuaternion()
    orientation = [q.GetReal()] + list(q.GetImaginary())
    t_mat = (
        Pose.from_list(list(translation) + orientation, TensorDeviceType())
        .get_matrix()
        .view(4, 4)
        .cpu()
        .numpy()
    )

    return t_mat, scale


def get_transform(pose):
    position = Gf.Vec3d(pose[:3])
    quat = pose[3:]
    rotation = Gf.Rotation(Gf.Quatf(quat[0], quat[1:]))

    mat_pose = Gf.Matrix4d()

    mat_pose.SetTransform(rotation, position)
    return mat_pose


def get_position_quat(pose, use_float: bool = True):
    quat = pose[3:]

    if use_float:
        position = Gf.Vec3f(pose[:3])

        quat = Gf.Quatf(quat[0], quat[1:])

    else:
        position = Gf.Vec3d(pose[:3])

        quat = Gf.Quatd(quat[0], quat[1:])
    return position, quat


def set_geom_mesh_attrs(mesh_geom: UsdGeom.Mesh, obs: Mesh, timestep=None):
    verts, faces = obs.get_mesh_data()
    mesh_geom.CreatePointsAttr(verts)
    mesh_geom.CreateFaceVertexCountsAttr([3 for _ in range(len(faces))])
    mesh_geom.CreateFaceVertexIndicesAttr(np.ravel(faces).tolist())
    mesh_geom.CreateSubdivisionSchemeAttr().Set(UsdGeom.Tokens.none)

    if obs.vertex_colors is not None:
        primvarsapi = UsdGeom.PrimvarsAPI(mesh_geom)
        primvar = primvarsapi.CreatePrimvar(
            "displayColor", Sdf.ValueTypeNames.Color3f, interpolation="faceVarying"
        )
        scale = 1.0
        # color needs to be in range of 0-1. Hence converting if the color is in [0,255]
        if max(np.ravel(obs.vertex_colors) > 1.0):
            scale = 255.0
        primvar.Set([Gf.Vec3f(x[0] / scale, x[1] / scale, x[2] / scale) for x in obs.vertex_colors])

    # low = np.min(verts, axis=0)
    # high = np.max(verts, axis=0)
    # mesh_geom.CreateExtentAttr([low, high])
    pose = obs.pose
    position = Gf.Vec3d(pose[:3])
    quat = pose[3:]
    q = Gf.Quatf(quat[0], quat[1:])

    # rotation = Gf.Rotation(Gf.Quatf(quat[0], quat[1:]))

    # mat_pose = Gf.Matrix4d()
    # mat_pose.SetTransform(rotation, position)
    # size = 1.0
    # mesh_geom.CreateSizeAttr(size)
    if timestep is not None:
        # UsdGeom.Xformable(mesh_geom).AddTransformOp().Set(time=timestep, value=mat_pose)
        a = UsdGeom.Xformable(mesh_geom)  #
        a.AddTranslateOp().Set(time=timestep, value=position)
        a.AddOrientOp().Set(time=timestep, value=q)
    else:
        a = UsdGeom.Xformable(mesh_geom)  #
        a.AddTranslateOp().Set(position)
        a.AddOrientOp().Set(q)

        # UsdGeom.Xformable(mesh_geom).AddTransformOp().Set(mat_pose)


def set_geom_cube_attrs(
    cube_geom: UsdGeom.Cube, dims: List[float], pose: List[float], timestep=None
):
    dims = Gf.Vec3d(np.ravel(dims).tolist())
    position = Gf.Vec3d(pose[:3])
    quat = pose[3:]
    q = Gf.Quatf(quat[0], quat[1:])
    # rotation = Gf.Rotation(q)

    # mat_pose = Gf.Matrix4d()
    # mat_pose.SetTransform(rotation, position)
    # mat = mat_pose
    # mat_scale = Gf.Matrix4d()
    # mat_scale.SetScale(dims)
    # mat = mat_scale * mat_pose
    size = 1.0
    cube_geom.CreateSizeAttr(size)
    # create scale:

    a = UsdGeom.Xformable(cube_geom)  #
    a.AddTranslateOp().Set(position)
    a.AddOrientOp().Set(q)

    # a.AddTransformOp().Set(mat)
    # scale will set the length to the given value
    a.AddScaleOp().Set(dims)


def set_geom_cylinder_attrs(
    cube_geom: UsdGeom.Cylinder, radius, height, pose: List[float], timestep=None
):
    position = Gf.Vec3d(pose[:3])
    quat = pose[3:]
    q = Gf.Quatf(quat[0], quat[1:])

    # create scale:
    cube_geom.CreateRadiusAttr(radius)
    cube_geom.CreateHeightAttr(height)
    a = UsdGeom.Xformable(cube_geom)  #
    a.AddTranslateOp().Set(position)
    a.AddOrientOp().Set(q)


def set_geom_sphere_attrs(
    sphere_geom: UsdGeom.Sphere, radius: float, pose: List[float], timestep=None
):
    position = Gf.Vec3d(pose[:3])
    quat = pose[3:]
    q = Gf.Quatf(quat[0], quat[1:])

    a = UsdGeom.Xformable(sphere_geom)  #
    a.AddTranslateOp().Set(position)
    a.AddOrientOp().Set(q)

    sphere_geom.CreateRadiusAttr(float(radius))


def set_cylinder_attrs(prim: UsdGeom.Cylinder, radius: float, height: float, pose, color=[]):
    # set size to 1:
    position = Gf.Vec3d(np.ravel(pose.xyz).tolist())
    quat = pose.so3.wxyz
    rotation = Gf.Rotation(Gf.Quatf(quat[0], quat[1:]))

    mat_pose = Gf.Matrix4d()
    mat_pose.SetTransform(rotation, position)

    prim.GetAttribute("height").Set(height)
    prim.GetAttribute("radius").Set(radius)

    UsdGeom.Xformable(prim).AddTransformOp().Set(mat_pose)


def get_cylinder_attrs(prim, cache=None, transform=None) -> Cylinder:
    size = prim.GetAttribute("size").Get()
    if size is None:
        size = 1.0
    height = prim.GetAttribute("height").Get() * size
    radius = prim.GetAttribute("radius").Get() * size

    mat, t_scale = get_prim_world_pose(cache, prim)

    if transform is not None:
        mat = transform @ mat
    # compute position and orientation on cuda:
    tensor_mat = torch.as_tensor(mat, device=torch.device("cuda", 0))
    pose = Pose.from_matrix(tensor_mat).tolist()
    return Cylinder(name=str(prim.GetPath()), pose=pose, height=height, radius=radius)


def get_capsule_attrs(prim, cache=None, transform=None) -> Cylinder:
    size = prim.GetAttribute("size").Get()
    if size is None:
        size = 1.0
    height = prim.GetAttribute("height").Get() * size
    radius = prim.GetAttribute("radius").Get() * size

    mat, t_scale = get_prim_world_pose(cache, prim)
    base = [0, 0, -height / 2]
    tip = [0, 0, height / 2]

    if transform is not None:
        mat = transform @ mat
    # compute position and orientation on cuda:
    tensor_mat = torch.as_tensor(mat, device=torch.device("cuda", 0))
    pose = Pose.from_matrix(tensor_mat).tolist()
    return Capsule(name=str(prim.GetPath()), pose=pose, base=base, tip=tip, radius=radius)


def get_cube_attrs(prim, cache=None, transform=None) -> Cuboid:
    # read cube size:
    size = prim.GetAttribute("size").Get()
    if size is None:
        size = 1.0
    dims = list(prim.GetAttribute("xformOp:scale").Get())
    # scale is 0.5 -> length of 1 will become 0.5,
    dims = [d * size for d in dims]
    if any([x <= 0 for x in dims]):
        raise ValueError("Negative or zero dimension")
    # dims = [x*2 for x in dims]
    mat, t_scale = get_prim_world_pose(cache, prim)

    if transform is not None:
        mat = transform @ mat
    # compute position and orientation on cuda:
    tensor_mat = torch.as_tensor(mat, device=torch.device("cuda", 0))
    pose = Pose.from_matrix(tensor_mat).tolist()
    return Cuboid(name=str(prim.GetPath()), pose=pose, dims=dims)


def get_sphere_attrs(prim, cache=None, transform=None) -> Sphere:
    # read cube information
    # scale = prim.GetAttribute("size").Get()
    size = prim.GetAttribute("size").Get()
    if size is None:
        size = 1.0
    radius = prim.GetAttribute("radius").Get()
    scale = prim.GetAttribute("xformOp:scale").Get()
    if scale is not None:
        radius = radius * max(list(scale)) * size

    if radius <= 0:
        raise ValueError("Negative or zero radius")
    # dims = [x*2 for x in dims]
    mat, t_scale = get_prim_world_pose(cache, prim)
    # position = list(prim.GetAttribute("xformOp:translate").Get())
    # q = prim.GetAttribute("xformOp:orient").Get()
    # orientation = [q.GetReal()] + list(q.GetImaginary())

    if transform is not None:
        mat = transform @ mat
    # compute position and orientation on cuda:
    tensor_mat = torch.as_tensor(mat, device=torch.device("cuda", 0))
    pose = Pose.from_matrix(tensor_mat).tolist()

    return Sphere(name=str(prim.GetPath()), pose=pose, radius=radius, position=pose[:3])


def get_mesh_attrs(prim, cache=None, transform=None) -> Mesh:
    # read cube information
    # scale = prim.GetAttribute("size").Get()
    points = list(prim.GetAttribute("points").Get())
    points = [np.ravel(x) for x in points]
    # points = np.ndarray(points)

    faces = list(prim.GetAttribute("faceVertexIndices").Get())

    face_count = list(prim.GetAttribute("faceVertexCounts").Get())
    # assume faces are 3:
    if len(faces) / 3 != len(face_count):
        log_warn(
            "Mesh faces "
            + str(len(faces) / 3)
            + " are not matching faceVertexCounts "
            + str(len(face_count))
        )
        return None
    faces = np.array(faces).reshape(len(face_count), 3).tolist()
    if prim.GetAttribute("xformOp:scale").IsValid():
        scale = list(prim.GetAttribute("xformOp:scale").Get())
    else:
        scale = [1.0, 1.0, 1.0]
    size = prim.GetAttribute("size").Get()
    if size is None:
        size = 1
    scale = [s * size for s in scale]

    mat, t_scale = get_prim_world_pose(cache, prim)
    # also get any world scale:
    scale = t_scale
    # position = list(prim.GetAttribute("xformOp:translate").Get())
    # q = prim.GetAttribute("xformOp:orient").Get()
    # orientation = [q.GetReal()] + list(q.GetImaginary())

    if transform is not None:
        mat = transform @ mat
    # compute position and orientation on cuda:
    tensor_mat = torch.as_tensor(mat, device=torch.device("cuda", 0))
    pose = Pose.from_matrix(tensor_mat).tolist()

    #

    m = Mesh(
        name=str(prim.GetPath()),
        pose=pose,
        vertices=points,
        faces=faces,
        scale=scale,
    )
    # print(len(m.vertices), max(m.faces))

    return m


def create_stage(
    name: str = "curobo_stage.usd",
    base_frame: str = "/world",
):
    stage = Usd.Stage.CreateNew(name)
    UsdGeom.SetStageUpAxis(stage, "Z")
    UsdGeom.SetStageMetersPerUnit(stage, 1)
    UsdPhysics.SetStageKilogramsPerUnit(stage, 1)
    xform = stage.DefinePrim(base_frame, "Xform")
    stage.SetDefaultPrim(xform)
    return stage


class UsdHelper:
    def __init__(self, use_float=True) -> None:
        self.stage = None
        self.dt = None
        self._use_float = use_float
        self._xform_cache = UsdGeom.XformCache()

    def create_stage(
        self,
        name: str = "curobo_stage.usd",
        base_frame: str = "/world",
        timesteps: Optional[int] = None,
        dt=0.02,
        interpolation_steps: float = 1,
    ):
        # print("name", name)
        self.stage = Usd.Stage.CreateNew(name)
        UsdGeom.SetStageUpAxis(self.stage, "Z")
        UsdGeom.SetStageMetersPerUnit(self.stage, 1)
        UsdPhysics.SetStageKilogramsPerUnit(self.stage, 1)
        xform = self.stage.DefinePrim(base_frame, "Xform")
        self.stage.SetDefaultPrim(xform)
        self.dt = dt
        self.interpolation_steps = interpolation_steps
        if timesteps is not None:
            self.stage.SetStartTimeCode(1)
            self.stage.SetEndTimeCode(timesteps * self.interpolation_steps)
            self.stage.SetTimeCodesPerSecond((1.0 / self.dt))
            # print(1.0 / self)

    def add_subroot(self, root="/world", sub_root="/obstacles", pose: Optional[Pose] = None):
        xform = self.stage.DefinePrim(join_path(root, sub_root), "Xform")
        if pose is not None:
            set_prim_transform(xform, pose.tolist(), use_float=self._use_float)

    def load_stage_from_file(self, file_path: str):
        self.stage = Usd.Stage.Open(file_path)

    def load_stage(self, stage: Usd.Stage):
        self.stage = stage

    def get_pose(self, prim_path: str, timecode: float = 0.0, inverse: bool = False) -> np.matrix:
        self._xform_cache.SetTime(timecode)
        reference_prim = self.stage.GetPrimAtPath(prim_path)
        r_T_w, _ = get_prim_world_pose(self._xform_cache, reference_prim, inverse=inverse)
        return r_T_w

    def get_obstacles_from_stage(
        self,
        only_paths: Optional[List[str]] = None,
        ignore_paths: Optional[List[str]] = None,
        only_substring: Optional[List[str]] = None,
        ignore_substring: Optional[List[str]] = None,
        reference_prim_path: Optional[str] = None,
        timecode: float = 0,
    ) -> WorldConfig:
        # read obstacles from usd by iterating through all prims:
        obstacles = {"cuboid": [], "sphere": None, "mesh": None, "cylinder": None, "capsule": None}
        r_T_w = None
        self._xform_cache.Clear()
        self._xform_cache.SetTime(timecode)
        if reference_prim_path is not None:
            reference_prim = self.stage.GetPrimAtPath(reference_prim_path)
            r_T_w, _ = get_prim_world_pose(self._xform_cache, reference_prim, inverse=True)
        all_items = self.stage.Traverse()
        for x in all_items:
            if only_paths is not None:
                if not any([str(x.GetPath()).startswith(k) for k in only_paths]):
                    continue
            if ignore_paths is not None:
                if any([str(x.GetPath()).startswith(k) for k in ignore_paths]):
                    continue
            if only_substring is not None:
                if not any([k in str(x.GetPath()) for k in only_substring]):
                    continue
            if ignore_substring is not None:
                if any([k in str(x.GetPath()) for k in ignore_substring]):
                    continue
            if x.IsA(UsdGeom.Cube):
                if obstacles["cuboid"] is None:
                    obstacles["cuboid"] = []
                cube = get_cube_attrs(x, cache=self._xform_cache, transform=r_T_w)
                obstacles["cuboid"].append(cube)
            elif x.IsA(UsdGeom.Sphere):
                if obstacles["sphere"] is None:
                    obstacles["sphere"] = []
                obstacles["sphere"].append(
                    get_sphere_attrs(x, cache=self._xform_cache, transform=r_T_w)
                )
            elif x.IsA(UsdGeom.Mesh):
                if obstacles["mesh"] is None:
                    obstacles["mesh"] = []
                m_data = get_mesh_attrs(x, cache=self._xform_cache, transform=r_T_w)
                if m_data is not None:
                    obstacles["mesh"].append(m_data)
            elif x.IsA(UsdGeom.Cylinder):
                if obstacles["cylinder"] is None:
                    obstacles["cylinder"] = []
                cube = get_cylinder_attrs(x, cache=self._xform_cache, transform=r_T_w)
                obstacles["cylinder"].append(cube)
            elif x.IsA(UsdGeom.Capsule):
                if obstacles["capsule"] is None:
                    obstacles["capsule"] = []
                cap = get_capsule_attrs(x, cache=self._xform_cache, transform=r_T_w)
                obstacles["capsule"].append(cap)
        world_model = WorldConfig(**obstacles)
        return world_model

    def add_world_to_stage(
        self,
        obstacles: WorldConfig,
        base_frame: str = "/world",
        obstacles_frame: str = "obstacles",
        base_t_obstacle_pose: Optional[Pose] = None,
        timestep: Optional[float] = None,
    ):
        # iterate through every obstacle type and create prims:

        self.add_subroot(base_frame, obstacles_frame, base_t_obstacle_pose)
        full_path = join_path(base_frame, obstacles_frame)
        prim_path = [
            self.get_prim_from_obstacle(o, full_path, timestep=timestep) for o in obstacles.objects
        ]
        return prim_path

    def get_prim_from_obstacle(
        self, obstacle: Obstacle, base_frame: str = "/world/obstacles", timestep=None
    ):

        if isinstance(obstacle, Cuboid):
            return self.add_cuboid_to_stage(obstacle, base_frame, timestep=timestep)
        elif isinstance(obstacle, Mesh):
            return self.add_mesh_to_stage(obstacle, base_frame, timestep=timestep)
        elif isinstance(obstacle, Sphere):
            return self.add_sphere_to_stage(obstacle, base_frame, timestep=timestep)
        elif isinstance(obstacle, Cylinder):
            return self.add_cylinder_to_stage(obstacle, base_frame, timestep=timestep)

        else:
            raise NotImplementedError

    def add_cuboid_to_stage(
        self,
        obstacle: Cuboid,
        base_frame: str = "/world/obstacles",
        timestep=None,
        enable_physics: bool = False,
    ):
        root_path = join_path(base_frame, obstacle.name)
        obj_geom = UsdGeom.Cube.Define(self.stage, root_path)
        obj_prim = self.stage.GetPrimAtPath(root_path)

        set_geom_cube_attrs(obj_geom, obstacle.dims, obstacle.pose, timestep=timestep)
        obj_prim.CreateAttribute("physics:rigidBodyEnabled", Sdf.ValueTypeNames.Bool, custom=False)
        obj_prim.GetAttribute("physics:rigidBodyEnabled").Set(enable_physics)

        if obstacle.color is not None:
            self.add_material(
                "material_" + obstacle.name, root_path, obstacle.color, obj_prim, obstacle.material
            )
        return root_path

    def add_cylinder_to_stage(
        self,
        obstacle: Cylinder,
        base_frame: str = "/world/obstacles",
        timestep=None,
        enable_physics: bool = False,
    ):
        root_path = join_path(base_frame, obstacle.name)
        obj_geom = UsdGeom.Cylinder.Define(self.stage, root_path)
        obj_prim = self.stage.GetPrimAtPath(root_path)

        set_geom_cylinder_attrs(
            obj_geom, obstacle.radius, obstacle.height, obstacle.pose, timestep=timestep
        )
        obj_prim.CreateAttribute("physics:rigidBodyEnabled", Sdf.ValueTypeNames.Bool, custom=False)
        obj_prim.GetAttribute("physics:rigidBodyEnabled").Set(enable_physics)

        if obstacle.color is not None:
            self.add_material(
                "material_" + obstacle.name, root_path, obstacle.color, obj_prim, obstacle.material
            )
        return root_path

    def add_sphere_to_stage(
        self,
        obstacle: Sphere,
        base_frame: str = "/world/obstacles",
        timestep=None,
        enable_physics: bool = False,
    ):
        root_path = join_path(base_frame, obstacle.name)
        obj_geom = UsdGeom.Sphere.Define(self.stage, root_path)
        obj_prim = self.stage.GetPrimAtPath(root_path)
        if obstacle.pose is None:
            obstacle.pose = obstacle.position + [1, 0, 0, 0]
        set_geom_sphere_attrs(obj_geom, obstacle.radius, obstacle.pose, timestep=timestep)
        obj_prim.CreateAttribute("physics:rigidBodyEnabled", Sdf.ValueTypeNames.Bool, custom=False)
        obj_prim.GetAttribute("physics:rigidBodyEnabled").Set(enable_physics)

        if obstacle.color is not None:
            self.add_material(
                "material_" + obstacle.name, root_path, obstacle.color, obj_prim, obstacle.material
            )
        return root_path

    def add_mesh_to_stage(
        self,
        obstacle: Mesh,
        base_frame: str = "/world/obstacles",
        timestep=None,
        enable_physics: bool = False,
    ):
        root_path = join_path(base_frame, obstacle.name)
        obj_geom = UsdGeom.Mesh.Define(self.stage, root_path)
        obj_prim = self.stage.GetPrimAtPath(root_path)
        # obstacle.update_material() # This does not get the correct materials
        set_geom_mesh_attrs(obj_geom, obstacle, timestep=timestep)

        obj_prim.CreateAttribute("physics:rigidBodyEnabled", Sdf.ValueTypeNames.Bool, custom=False)
        obj_prim.GetAttribute("physics:rigidBodyEnabled").Set(enable_physics)

        if obstacle.color is not None:
            self.add_material(
                "material_" + obstacle.name, root_path, obstacle.color, obj_prim, obstacle.material
            )

        return root_path

    def get_obstacle_from_prim(self, prim_path: str) -> Obstacle:
        pass

    def write_stage_to_file(self, file_path: str, flatten: bool = False):
        if flatten:
            usd_str = self.stage.Flatten().ExportToString()
        else:
            usd_str = self.stage.GetRootLayer().ExportToString()
        with open(file_path, "w") as f:
            f.write(usd_str)

    def create_animation(
        self,
        robot_world_cfg: WorldConfig,
        pose: Pose,
        base_frame="/world",
        robot_frame="/robot",
        dt: float = 0.02,
    ):
        """Create animation, given meshes and pose

        Args:
            prim_names: _description_
            pose: [ timesteps, n_meshes, pose]
            dt: _description_. Defaults to 0.02.
        """
        prim_names = self.add_world_to_stage(
            robot_world_cfg, base_frame=base_frame, obstacles_frame=robot_frame, timestep=0
        )
        for i, i_val in enumerate(prim_names):
            curr_prim = self.stage.GetPrimAtPath(i_val)
            form = UsdGeom.Xformable(curr_prim).GetOrderedXformOps()
            if len(form) < 2:
                log_warn("Pose transformation not found" + i_val)
                continue

            pos_form = form[0]
            quat_form = form[1]
            use_float = True  # default is float
            for t in range(pose.batch):
                c_p, c_q = get_position_quat(pose.get_index(t, i).tolist(), use_float)
                pos_form.Set(time=t * self.interpolation_steps, value=c_p)
                quat_form.Set(time=t * self.interpolation_steps, value=c_q)
                # c_t = get_transform(pose.get_index(t, i).tolist())
                # form.Set(time=t * self.interpolation_steps, value=c_t)

    def create_obstacle_animation(
        self,
        obstacles: List[List[Obstacle]],
        base_frame: str = "/world",
        obstacles_frame: str = "robot_base",
    ):
        # add obstacles to stage:
        prim_paths = self.add_world_to_stage(
            WorldConfig(objects=obstacles[0]),
            base_frame=base_frame,
            obstacles_frame=obstacles_frame,
        )

        #
        for t in range(len(obstacles)):
            current_obs = obstacles[t]
            for j in range(len(current_obs)):
                obs = current_obs[j]
                obs_name = join_path(join_path(base_frame, obstacles_frame), obs.name)
                if obs_name not in prim_paths:
                    log_warn("Obstacle not found")
                    continue
                #
                prim = self.stage.GetPrimAtPath(obs_name)
                form = UsdGeom.Xformable(prim).GetOrderedXformOps()

                pos_form = form[0]
                c_p = Gf.Vec3d(obs.position)
                pos_form.Set(time=t * self.interpolation_steps, value=c_p)

    def create_linkpose_robot_animation(
        self,
        robot_usd_path: str,
        link_names: List[str],
        joint_names: List[str],
        pose: Pose,
        robot_base_frame="/world/robot",
        local_asset_path="assets/",
        write_robot_usd_path="assets/",
        robot_asset_prim_path="/panda",
    ):
        """Create animation, given meshes and pose

        Args:
            prim_names: _description_
            pose: [ timesteps, n_meshes, pose]
            dt: _description_. Defaults to 0.02.
        """
        link_prims, joint_prims = self.load_robot_usd(
            robot_usd_path,
            link_names,
            joint_names,
            robot_base_frame=robot_base_frame,
            local_asset_path=local_asset_path,
            write_asset_path=write_robot_usd_path,
            robot_asset_prim_path=robot_asset_prim_path,
        )

        for i, i_val in enumerate(link_names):
            if i_val not in link_prims:
                log_warn("Link not found in usd: " + i_val)
                continue
            form = UsdGeom.Xformable(link_prims[i_val]).GetOrderedXformOps()
            if len(form) < 2:
                log_warn("Pose transformation not found" + i_val)
                continue

            pos_form = form[0]
            quat_form = form[1]
            use_float = False
            if link_prims[i_val].GetAttribute("xformOp:orient").IsValid():
                if isinstance(link_prims[i_val].GetAttribute("xformOp:orient").Get(), Gf.Quatf):
                    use_float = True

            for t in range(pose.batch):
                c_p, c_q = get_position_quat(pose.get_index(t, i).tolist(), use_float)
                pos_form.Set(time=t * self.interpolation_steps, value=c_p)
                quat_form.Set(time=t * self.interpolation_steps, value=c_q)

    def add_material(
        self,
        material_name: str,
        object_path: str,
        color: List[float],
        obj_prim: Usd.Prim,
        material: Material = Material(),
    ):
        mat_path = join_path(object_path, material_name)
        material_usd = UsdShade.Material.Define(self.stage, mat_path)
        pbrShader = UsdShade.Shader.Define(self.stage, join_path(mat_path, "PbrShader"))
        pbrShader.CreateIdAttr("UsdPreviewSurface")
        pbrShader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(material.roughness)
        pbrShader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(material.metallic)
        pbrShader.CreateInput("specularColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(color[:3]))
        pbrShader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(color[:3]))
        pbrShader.CreateInput("baseColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(color[:3]))

        pbrShader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(color[3])
        material_usd.CreateSurfaceOutput().ConnectToSource(pbrShader.ConnectableAPI(), "surface")
        obj_prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
        UsdShade.MaterialBindingAPI(obj_prim).Bind(material_usd)
        return material_usd

    def save(self):
        self.stage.Save()

    @staticmethod
    def write_trajectory_animation(
        robot_model_file: str,
        world_model: WorldConfig,
        q_start: JointState,
        q_traj: JointState,
        dt: float = 0.02,
        save_path: str = "out.usd",
        tensor_args: TensorDeviceType = TensorDeviceType(),
        interpolation_steps: float = 1.0,
        robot_base_frame="robot",
        base_frame="/world",
        kin_model: Optional[CudaRobotModel] = None,
        visualize_robot_spheres: bool = True,
        robot_color: Optional[List[float]] = None,
        flatten_usd: bool = False,
        goal_pose: Optional[Pose] = None,
        goal_color: Optional[List[float]] = None,
    ):
        if kin_model is None:
            config_file = load_yaml(join_path(get_robot_configs_path(), robot_model_file))
            if "robot_cfg" not in config_file:
                config_file["robot_cfg"] = config_file
            config_file["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
            robot_cfg = CudaRobotModelConfig.from_data_dict(
                config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args
            )
            kin_model = CudaRobotModel(robot_cfg)
        m = kin_model.get_robot_link_meshes()
        offsets = [x.pose for x in m]
        robot_mesh_model = WorldConfig(mesh=m)
        if robot_color is not None:
            robot_mesh_model.add_color(robot_color)
            robot_mesh_model.add_material(Material(metallic=0.4))
        if goal_pose is not None:
            kin_model.link_names
            if kin_model.ee_link in kin_model.kinematics_config.mesh_link_names:
                index = kin_model.kinematics_config.mesh_link_names.index(kin_model.ee_link)
                gripper_mesh = m[index]
            if len(goal_pose.shape) == 1:
                goal_pose = goal_pose.unsqueeze(0)
            if len(goal_pose.shape) == 2:
                goal_pose = goal_pose.unsqueeze(0)
            for i in range(goal_pose.n_goalset):
                g = goal_pose.get_index(0, i).to_list()
                world_model.add_obstacle(
                    Mesh(
                        file_path=gripper_mesh.file_path,
                        pose=g,
                        name="goal_idx_" + str(i),
                        color=goal_color,
                    )
                )
        usd_helper = UsdHelper()
        usd_helper.create_stage(
            save_path,
            timesteps=q_traj.position.shape[0],
            dt=dt,
            interpolation_steps=interpolation_steps,
            base_frame=base_frame,
        )
        if world_model is not None:
            usd_helper.add_world_to_stage(world_model, base_frame=base_frame)

        animation_links = kin_model.kinematics_config.mesh_link_names
        animation_poses = kin_model.get_link_poses(q_traj.position.contiguous(), animation_links)
        # add offsets for visual mesh:
        for i, ival in enumerate(offsets):
            offset_pose = Pose.from_list(ival)
            new_pose = Pose(
                animation_poses.position[:, i, :], animation_poses.quaternion[:, i, :]
            ).multiply(offset_pose)
            animation_poses.position[:, i, :] = new_pose.position
            animation_poses.quaternion[:, i, :] = new_pose.quaternion

        robot_base_frame = join_path(base_frame, robot_base_frame)

        usd_helper.create_animation(
            robot_mesh_model, animation_poses, base_frame, robot_frame=robot_base_frame
        )
        if visualize_robot_spheres:
            # visualize robot spheres:
            sphere_traj = kin_model.get_robot_as_spheres(q_traj.position)
            # change color:
            for s in sphere_traj:
                for k in s:
                    k.color = [0, 0.27, 0.27, 1.0]
            usd_helper.create_obstacle_animation(
                sphere_traj, base_frame=base_frame, obstacles_frame="curobo/robot_collision"
            )
        usd_helper.write_stage_to_file(save_path, flatten=flatten_usd)

    @staticmethod
    def load_robot(
        robot_model_file: str,
        tensor_args: TensorDeviceType = TensorDeviceType(),
    ) -> CudaRobotModel:
        config_file = load_yaml(join_path(get_robot_configs_path(), robot_model_file))
        if "robot_cfg" in config_file:
            config_file = config_file["robot_cfg"]
        config_file["kinematics"]["load_link_names_with_mesh"] = True
        # config_file["robot_cfg"]["kinematics"]["use_usd_kinematics"] = False

        robot_cfg = CudaRobotModelConfig.from_data_dict(
            config_file["kinematics"], tensor_args=tensor_args
        )

        kin_model = CudaRobotModel(robot_cfg)
        return kin_model

    @staticmethod
    def write_trajectory_animation_with_robot_usd(
        robot_model_file: str,
        world_model: Union[WorldConfig, None],
        q_start: JointState,
        q_traj: JointState,
        dt: float = 0.02,
        save_path: str = "out.usd",
        tensor_args: TensorDeviceType = TensorDeviceType(),
        interpolation_steps: float = 1.0,
        write_robot_usd_path: str = "assets/",
        robot_base_frame: str = "robot",
        robot_usd_local_reference: str = "assets/",
        base_frame="/world",
        kin_model: Optional[CudaRobotModel] = None,
        visualize_robot_spheres: bool = True,
        robot_asset_prim_path=None,
        robot_color: Optional[List[float]] = None,
        flatten_usd: bool = False,
        goal_pose: Optional[Pose] = None,
        goal_color: Optional[List[float]] = None,
    ):
        usd_exists = False
        # if usd file doesn't exist, fall back to urdf animation script

        if kin_model is None:
            robot_model_file = load_yaml(join_path(get_robot_configs_path(), robot_model_file))
            if "robot_cfg" in robot_model_file:
                robot_model_file = robot_model_file["robot_cfg"]
            robot_model_file["kinematics"]["load_link_names_with_mesh"] = True
            robot_model_file["kinematics"]["use_usd_kinematics"] = True
            if "usd_path" in robot_model_file["kinematics"]:

                usd_exists = file_exists(
                    join_path(get_assets_path(), robot_model_file["kinematics"]["usd_path"])
                )
            else:
                usd_exists = False
        else:
            if kin_model.generator_config.usd_path is not None:
                usd_exists = file_exists(kin_model.generator_config.usd_path)
        if robot_color is not None:
            log_warn(
                "robot_color is not supported when using robot from usd, "
                + "using urdf mode instead to write usd file"
            )
            usd_exists = False
        if not usd_exists:
            log_info("robot usd not found, using urdf animation instead")
            robot_model_file["kinematics"]["use_usd_kinematics"] = False
            return UsdHelper.write_trajectory_animation(
                robot_model_file,
                world_model,
                q_start,
                q_traj,
                dt,
                save_path,
                tensor_args,
                interpolation_steps,
                robot_base_frame=robot_base_frame,
                base_frame=base_frame,
                kin_model=kin_model,
                visualize_robot_spheres=visualize_robot_spheres,
                robot_color=robot_color,
                flatten_usd=flatten_usd,
                goal_pose=goal_pose,
                goal_color=goal_color,
            )
        if kin_model is None:
            robot_cfg = CudaRobotModelConfig.from_data_dict(
                robot_model_file["kinematics"], tensor_args=tensor_args
            )

            kin_model = CudaRobotModel(robot_cfg)

        if robot_asset_prim_path is None:
            robot_asset_prim_path = kin_model.kinematics_parser.robot_prim_root

        robot_base_frame = join_path(base_frame, robot_base_frame)

        robot_usd_path = kin_model.generator_config.usd_path

        usd_helper = UsdHelper()
        usd_helper.create_stage(
            save_path,
            timesteps=q_traj.position.shape[0],
            dt=dt,
            interpolation_steps=interpolation_steps,
            base_frame=base_frame,
        )
        if world_model is not None:
            usd_helper.add_world_to_stage(world_model, base_frame=base_frame)

        animation_links = kin_model.kinematics_config.mesh_link_names
        animation_poses = kin_model.get_link_poses(q_traj.position, animation_links)

        usd_helper.create_linkpose_robot_animation(
            robot_usd_path,
            animation_links,
            kin_model.joint_names,
            animation_poses,
            local_asset_path=robot_usd_local_reference,
            write_robot_usd_path=write_robot_usd_path,
            robot_base_frame=robot_base_frame,
            robot_asset_prim_path=robot_asset_prim_path,
        )
        if visualize_robot_spheres:
            # visualize robot spheres:
            sphere_traj = kin_model.get_robot_as_spheres(q_traj.position)
            # change color:
            for s in sphere_traj:
                for k in s:
                    k.color = [0, 0.27, 0.27, 1.0]
            usd_helper.create_obstacle_animation(
                sphere_traj, base_frame=base_frame, obstacles_frame="curobo/robot_collision"
            )
        usd_helper.write_stage_to_file(save_path, flatten=flatten_usd)

    @staticmethod
    def create_grid_usd(
        usds_path: Union[str, List[str]],
        save_path: str,
        base_frame: str,
        max_envs: int,
        max_timecode: float,
        x_space: float,
        y_space: float,
        x_per_row: int,
        local_asset_path: str,
        dt: float = 0.02,
        interpolation_steps: int = 1,
        prefix_string: Optional[str] = None,
        flatten_usd: bool = False,
    ):
        # create stage:
        usd_helper = UsdHelper()
        usd_helper.create_stage(
            save_path,
            timesteps=max_timecode,
            dt=dt,
            interpolation_steps=interpolation_steps,
            base_frame=base_frame,
        )

        # read all usds:
        if isinstance(usds_path, list):
            files = usds_path
        else:
            files = get_files_from_dir(usds_path, [".usda", ".usd"], prefix_string)
        # get count and clamp to max:
        n_envs = min(len(files), max_envs)
        # create grid
        # :
        count_x = x_per_row
        count_y = int(np.ceil((n_envs) / x_per_row))
        x_set = np.linspace(0, x_space * count_x, count_x)
        y_set = np.linspace(0, y_space * count_y, count_y)
        xv, yv = np.meshgrid(x_set, y_set)
        xv = np.ravel(xv)
        yv = np.ravel(yv)

        # define prim + add reference:

        for i in range(n_envs):
            world_usd_path = files[i]
            env_base_frame = (
                base_frame + "/grid_" + get_filename(world_usd_path, remove_extension=True)
            )
            prim = usd_helper.stage.DefinePrim(env_base_frame, "Xform")
            set_prim_transform(prim, [xv[i], yv[i], 0, 1, 0, 0, 0])
            ref = prim.GetReferences()
            ref.AddReference(assetPath=join_path(local_asset_path, get_filename(world_usd_path)))

        # write usd to disk:

        usd_helper.write_stage_to_file(save_path, flatten=flatten_usd)

    def load_robot_usd(
        self,
        robot_usd_path: str,
        link_names: List[str],
        joint_names: List[str],
        robot_base_frame="/world/robot",
        write_asset_path="assets/",
        local_asset_path="assets/",
        robot_asset_prim_path="/panda",
    ):
        # copy robot prim and it's derivatives to a seperate usd:
        robot_usd_name = get_filename(robot_usd_path)

        out_path = join_path(write_asset_path, robot_usd_name)
        out_local_path = join_path(local_asset_path, robot_usd_name)
        if not file_exists(out_path) or not file_exists(out_local_path):
            robot_stage = Usd.Stage.Open(robot_usd_path)  # .Flatten()  # .Flatten()
            # set pose to zero for root prim:

            prim = robot_stage.GetPrimAtPath(robot_asset_prim_path)
            if not prim.IsValid():
                log_error(
                    "robot prim is not valid : " + robot_asset_prim_path + " " + robot_usd_path
                )
            set_prim_transform(prim, [0, 0, 0, 1, 0, 0, 0])
            robot_stage = robot_stage.Flatten()
            robot_stage.Export(out_path)
            robot_stage.Export(out_local_path)

        # create a base prim:
        prim = self.stage.DefinePrim(robot_base_frame)
        ref = prim.GetReferences()

        ref.AddReference(
            assetPath=join_path(local_asset_path, robot_usd_name), primPath=robot_asset_prim_path
        )
        link_prims, joint_prims = self.get_robot_prims(link_names, joint_names, robot_base_frame)
        return link_prims, joint_prims

    def get_robot_prims(
        self, link_names: List[str], joint_names: List[str], robot_base_path: str = "/world/robot"
    ):
        all_prims = [x for x in self.stage.Traverse()]
        joint_prims = {}
        link_prims = {}
        for j_idx, j in enumerate(joint_names):
            for k in range(len(all_prims)):
                current_prim = all_prims[k]
                prim_path = current_prim.GetPath().pathString
                if robot_base_path in prim_path and j in prim_path:
                    joint_prims[j] = current_prim
                    current_prim.GetAttribute("physics:jointEnabled").Set(False)
        for j_idx, j in enumerate(link_names):
            for k in range(len(all_prims)):
                current_prim = all_prims[k]
                prim_path = current_prim.GetPath().pathString
                if (
                    robot_base_path in prim_path
                    and j in prim_path
                    and "geometry" not in prim_path
                    and "joint" not in prim_path
                    and current_prim.GetTypeName() == "Xform"
                ):
                    link_prims[j] = current_prim

                    # stat = current_prim.GetAttribute("physics:rigidBodyEnabled")
                    current_prim.GetAttribute("physics:rigidBodyEnabled").Set(False)

        return link_prims, joint_prims

    def update_robot_joint_state(self, joint_prims: List[Usd.Prim], js: JointState, timestep: int):
        for j_idx, j in enumerate(js.joint_names):
            if timestep is not None:
                joint_prims[j].GetAttribute("drive:angular:physics:targetPosition").Set(
                    time=timestep, value=np.degrees(js.position[..., j_idx].item())
                )
            else:
                joint_prims[j].GetAttribute("drive:angular:physics:targetPosition").Set(
                    value=np.degrees(js.position[..., j_idx].item())
                )

    @staticmethod
    def write_motion_gen_log(
        result: MotionGenResult,
        robot_file: Union[str, RobotConfig],
        world_config: Union[None, WorldConfig],
        start_state: JointState,
        goal_pose: Pose,
        save_prefix: str = "log",
        write_ik: bool = False,
        write_trajopt: bool = False,
        write_graph: bool = False,
        goal_object: Optional[Obstacle] = None,
        overlay_ik: bool = False,
        overlay_trajopt: bool = False,
        visualize_robot_spheres: bool = True,
        link_spheres: Optional[torch.Tensor] = None,
        grid_space: float = 1.0,
        write_robot_usd_path="assets/",
        robot_asset_prim_path="/panda",
        fps: int = 24,
        link_poses: Optional[Dict[str, Pose]] = None,
        flatten_usd: bool = False,
    ):
        if goal_object is None:
            log_warn("Using franka gripper as goal object")
            goal_object = Mesh(
                name="target_gripper",
                file_path=join_path(
                    get_assets_path(),
                    "robot/franka_description/meshes/visual/hand.dae",
                ),
                color=[0.0, 0.8, 0.1, 1.0],
                pose=goal_pose.to_list(),
            )

        if goal_object is not None:
            goal_object.pose = np.ravel(goal_pose.tolist()).tolist()
            world_config = world_config.clone()
            world_config.add_obstacle(goal_object)
        if link_poses is not None:
            link_goals = []
            for k in link_poses.keys():
                link_goals.append(
                    Mesh(
                        name="target_" + k,
                        file_path=join_path(
                            get_assets_path(),
                            "robot/franka_description/meshes/visual/hand.dae",
                        ),
                        color=[0.0, 0.8, 0.1, 1.0],
                        pose=link_poses[k].to_list(),
                    )
                )
                world_config.add_obstacle(link_goals[-1])
        kin_model = UsdHelper.load_robot(robot_file)
        if link_spheres is not None:
            kin_model.kinematics_config.link_spheres = link_spheres
        if write_graph:
            log_error("Logging graph planner is not supported")
        if write_ik:
            x_space = y_space = grid_space
            if overlay_ik:
                x_space = y_space = 0.0

            ik_iter_steps = result.debug_info["ik_result"].debug_info["solver"]["steps"]
            # convert ik_iter to a trajectory:
            ik_iter_steps = torch.cat(
                [torch.cat(ik_iter_steps[i], dim=1) for i in range(len(ik_iter_steps))], dim=1
            )
            vis_traj = ik_iter_steps
            num_seeds, n_iters, dof = vis_traj.shape
            usd_paths = []
            for j in tqdm(range(num_seeds)):
                current_traj = vis_traj[j].view(-1, dof)[:, :]  # we remove last timestep

                usd_paths.append(save_prefix + "_ik_seed_" + str(j) + ".usd")
                UsdHelper.write_trajectory_animation_with_robot_usd(
                    robot_file,
                    world_config,
                    start_state,
                    JointState.from_position(current_traj),
                    dt=(1 / fps),
                    save_path=usd_paths[-1],
                    base_frame="/world_base_" + str(j),
                    kin_model=kin_model,
                    visualize_robot_spheres=visualize_robot_spheres,
                    write_robot_usd_path=write_robot_usd_path,
                    robot_asset_prim_path=robot_asset_prim_path,
                    flatten_usd=flatten_usd,
                )

            UsdHelper.create_grid_usd(
                usd_paths,
                save_prefix + "_grid_ik.usd",
                base_frame="/world",
                max_envs=len(usd_paths),
                max_timecode=n_iters,
                x_space=x_space,
                y_space=y_space,
                x_per_row=int(math.floor(math.sqrt(len(usd_paths)))),
                local_asset_path="",
                dt=(1.0 / fps),
            )
        if write_trajopt:
            if "trajopt_result" not in result.debug_info:
                log_warn("Trajopt result was not found in debug information")
                return
            trajectory_iter_steps = result.debug_info["trajopt_result"].debug_info["solver"][
                "steps"
            ]
            vis_traj = []
            for i in range(len(trajectory_iter_steps)):
                vis_traj += trajectory_iter_steps[i]

            full_traj = torch.cat(vis_traj, dim=0)
            num_seeds, h, dof = vis_traj[-1].shape
            n, _, _ = full_traj.shape  # this will have iterations
            full_traj = full_traj.view(-1, num_seeds, h, dof)
            n_steps = full_traj.shape[0]

            full_traj = full_traj.transpose(0, 1).contiguous()  # n_seeds, n_steps, h, dof
            n1, n2, _, _ = full_traj.shape

            full_traj = torch.cat(
                (start_state.position.view(1, 1, 1, -1).repeat(n1, n2, 1, 1), full_traj), dim=-2
            )
            usd_paths = []
            finetune_usd_paths = []
            for j in tqdm(range(num_seeds)):
                current_traj = full_traj[j].view(-1, dof)  # we remove last timestep
                # add start state to current trajectory since it's not in the optimization:
                usd_paths.append(save_prefix + "_trajopt_seed_" + str(j) + ".usd")
                UsdHelper.write_trajectory_animation_with_robot_usd(
                    robot_file,
                    world_config,
                    start_state,
                    JointState.from_position(current_traj),
                    dt=(1.0 / fps),
                    save_path=usd_paths[-1],
                    base_frame="/world_base_" + str(j),
                    kin_model=kin_model,
                    visualize_robot_spheres=visualize_robot_spheres,
                    write_robot_usd_path=write_robot_usd_path,
                    robot_asset_prim_path=robot_asset_prim_path,
                    flatten_usd=flatten_usd,
                )
            # add finetuning step:

            if "finetune_trajopt_result" in result.debug_info and True:
                finetune_iter_steps = result.debug_info["finetune_trajopt_result"].debug_info[
                    "solver"
                ]["steps"]

                vis_traj = []
                if finetune_iter_steps is not None:
                    for i in range(len(finetune_iter_steps)):
                        vis_traj += finetune_iter_steps[i]
                full_traj = torch.cat(vis_traj, dim=0)
                num_seeds, h, dof = vis_traj[-1].shape
                n, _, _ = full_traj.shape  # this will have iterations
                # print(full_traj.shape)
                full_traj = full_traj.view(-1, num_seeds, h, dof)
                n1, n2, _, _ = full_traj.shape

                full_traj = torch.cat(
                    (start_state.position.view(1, 1, 1, -1).repeat(n1, n2, 1, 1), full_traj), dim=-2
                )

                # n_steps = full_traj.shape[0]

                full_traj = full_traj.transpose(0, 1).contiguous()  # n_seeds, n_steps, h, dof
                for j in tqdm(range(num_seeds)):
                    current_traj = full_traj[j][-1].view(-1, dof)

                    finetune_usd_paths.append(save_prefix + "_finetune_seed_" + str(j) + ".usd")
                    UsdHelper.write_trajectory_animation_with_robot_usd(
                        robot_file,
                        world_config,
                        start_state,
                        JointState.from_position(current_traj),
                        dt=(1.0 / fps),
                        save_path=finetune_usd_paths[-1],
                        base_frame="/world_base_" + str(j),
                        kin_model=kin_model,
                        visualize_robot_spheres=visualize_robot_spheres,
                        write_robot_usd_path=write_robot_usd_path,
                        robot_asset_prim_path=robot_asset_prim_path,
                        flatten_usd=flatten_usd,
                    )
            x_space = y_space = grid_space
            if overlay_trajopt:
                x_space = y_space = 0.0
            if len(finetune_usd_paths) == 1:
                usd_paths.append(finetune_usd_paths[0])
            UsdHelper.create_grid_usd(
                usd_paths,
                save_prefix + "_grid_trajopt.usd",
                base_frame="/world",
                max_envs=len(usd_paths),
                max_timecode=n_steps * h,
                x_space=x_space,
                y_space=y_space,
                x_per_row=int(math.floor(math.sqrt(len(usd_paths)))),
                local_asset_path="",
                dt=(1.0 / fps),
            )
            if False and len(finetune_usd_paths) > 1:
                UsdHelper.create_grid_usd(
                    finetune_usd_paths,
                    save_prefix + "_grid_finetune_trajopt.usd",
                    base_frame="/world",
                    max_envs=len(finetune_usd_paths),
                    max_timecode=n_steps * h,
                    x_space=x_space,
                    y_space=y_space,
                    x_per_row=int(math.floor(math.sqrt(len(finetune_usd_paths)))),
                    local_asset_path="",
                    dt=(1.0 / fps),
                )
