import omni
from omni.isaac.core.robots import Robot
import omni.timeline
import carb
import asyncio

def print_robot_actions_info(robot_prim_path="/World/piper_description"):
    """打印机器人的所有动作和关节信息"""
    print("\n========== 机器人动作和关节信息 ==========")
    
    # 获取机器人对象
    try:
        robot = Robot(prim_path=robot_prim_path)
        robot.initialize()
        print(f"成功加载机器人: {robot_prim_path}")
    except Exception as e:
        print(f"无法加载机器人: {e}")
        return
    
    # 1. 打印关节数量
    num_joints = robot.num_dof
    print(f"\n关节数量: {num_joints}")
    
    # 2. 打印关节名称
    joint_names = robot.dof_names
    print("\n关节名称:")
    for i, name in enumerate(joint_names):
        print(f"  [{i}] {name}")
    
    # 3. 打印关节类型
    print("\n关节类型:")
    for i, name in enumerate(joint_names):
        # 大多数机器人关节是旋转关节
        if "prismatic" in name.lower():
            joint_type_str = "PRISMATIC"
        else:
            joint_type_str = "REVOLUTE"
        print(f"  [{i}] {name}: {joint_type_str}")
    
    # 4. 打印关节位置信息
    positions = robot.get_joint_positions()
    print("\n当前关节位置:")
    for i, pos in enumerate(positions):
        print(f"  [{i}] {joint_names[i]}: {pos:.4f}")
    
    # 5. 打印关节速度
    velocities = robot.get_joint_velocities()
    print("\n当前关节速度:")
    for i, vel in enumerate(velocities):
        print(f"  [{i}] {joint_names[i]}: {vel:.4f}")
    
    # 6. 打印关节力矩
    efforts = robot.get_applied_joint_efforts()
    print("\n当前关节力矩:")
    for i, eff in enumerate(efforts):
        print(f"  [{i}] {joint_names[i]}: {eff:.4f}")
    
    # 使用Isaac Sim 4.2的API获取关节限制
    print("\n关节限制 (使用Isaac Sim 4.2 API):")
    try:
        from omni.isaac.core.articulations import Articulation
        from pxr import UsdPhysics
        
        # 创建一个Articulation对象
        articulation = Articulation(prim_path=robot_prim_path)
        articulation.initialize()
        
        # 通过USD接口获取关节限制
        from pxr import Usd, UsdPhysics
        stage = omni.usd.get_context().get_stage()
        
        for i, joint_name in enumerate(joint_names):
            # 尝试从名称推断关节路径
            # 注意：这里假设关节路径格式为 robotPath/jointName
            joint_path = f"{robot_prim_path}/{joint_name}"
            joint_prim = stage.GetPrimAtPath(joint_path)
            
            if not joint_prim:
                # 尝试其他可能的路径格式
                possible_paths = [
                    f"{robot_prim_path}/joints/{joint_name}",
                    f"{robot_prim_path}/{joint_name}_joint"
                ]
                
                for path in possible_paths:
                    joint_prim = stage.GetPrimAtPath(path)
                    if joint_prim:
                        joint_path = path
                        break
            
            if joint_prim:
                joint_api = UsdPhysics.RevoluteJoint(joint_prim)
                if joint_api:
                    lower_attr = joint_api.GetLowerAttr()
                    upper_attr = joint_api.GetUpperAttr()
                    
                    if lower_attr and upper_attr:
                        lower = lower_attr.Get() if lower_attr.HasValue() else float('-inf')
                        upper = upper_attr.Get() if upper_attr.HasValue() else float('inf')
                        print(f"  [{i}] {joint_name}: 下限={lower:.4f}, 上限={upper:.4f}")
                        continue
            
            # 如果无法找到关节或获取限制，使用默认值
            print(f"  [{i}] {joint_name}: 下限=-π, 上限=π (默认值)")
    
    except Exception as e:
        print(f"  无法通过USD接口获取关节限制: {e}")
        print("  使用默认值: 下限=-π, 上限=π")
    
    # 8. 驱动类型和控制器参数
    print("\n驱动类型和控制器参数 (使用Action Graph节点获取):")
    print("  请使用下面的方法获取Action Graph中的控制器信息")
    
    print("\n========================================")


# 使用兼容的方式获取Action Graph信息
def print_action_graph_info_compatible():
    """使用兼容Isaac Sim 4.2的方式打印Action Graph信息"""
    print("\n========== Action Graph节点信息 (兼容方式) ==========")
    
    try:
        # 在Isaac Sim 4.2中，我们需要使用不同的API获取Action Graph
        import omni.kit
        from omni.kit.viewport.utility import get_viewport_from_window_name
        viewport = get_viewport_from_window_name("Viewport")
        
        # 打印Action Graph信息 - 方法1：使用USD查询
        import omni.usd
        from pxr import Usd, Sdf
        
        # 获取当前舞台
        stage = omni.usd.get_context().get_stage()
        
        # 查找所有ActionGraph节点
        print("\nAction Graph节点查找方法1 - USD路径查询:")
        action_graph_paths = []
        
        # 查找可能的Action Graph节点路径
        possible_paths = ["/World/ActionGraph", "/ActionGraph"]
        for path in possible_paths:
            prim = stage.GetPrimAtPath(path)
            if prim:
                action_graph_paths.append(path)
                print(f"  找到Action Graph路径: {path}")
                
                # 打印子节点
                children = prim.GetChildren()
                if children:
                    print("  子节点:")
                    for child in children:
                        print(f"    - {child.GetPath()}")
                        
                        # 尝试获取更多信息
                        attrs = child.GetAttributes()
                        if attrs:
                            print(f"    属性数量: {len(attrs)}")
                            # 打印前10个属性
                            for i, attr in enumerate(attrs):
                                if i >= 10:
                                    break
                                print(f"      {attr.GetName()}: {attr.Get() if attr.IsValid() else 'N/A'}")
        
        if not action_graph_paths:
            print("  无法在常见路径找到Action Graph节点")
            
            # 尝试搜索整个场景
            print("\n  尝试搜索整个场景中的ActionGraph相关节点:")
            
            def traverse_stage(prim):
                paths = []
                if "ActionGraph" in prim.GetName() or "articulation_controller" in prim.GetName().lower():
                    paths.append(prim.GetPath())
                
                for child in prim.GetChildren():
                    paths.extend(traverse_stage(child))
                return paths
            
            root = stage.GetPseudoRoot()
            all_action_paths = traverse_stage(root)
            
            if all_action_paths:
                print(f"  找到 {len(all_action_paths)} 个可能的Action Graph相关节点:")
                for path in all_action_paths:
                    print(f"    - {path}")
            else:
                print("  未找到任何ActionGraph相关节点")
                
        # 尝试查找ArticulationController节点
        print("\n尝试找到ArticulationController节点:")
        controller_found = False
        
        for prim in Usd.PrimRange(stage.GetPseudoRoot()):
            prim_type = prim.GetTypeName()
            prim_name = prim.GetName()
            
            if "ArticulationController" in str(prim_type) or "ArticulationController" in prim_name:
                print(f"  找到控制器: {prim.GetPath()}")
                controller_found = True
                
                # 打印属性
                attrs = prim.GetAttributes()
                if attrs:
                    print("  控制器属性:")
                    for attr in attrs:
                        if attr.IsValid():
                            print(f"    - {attr.GetName()}: {attr.Get() if attr.HasValue() else 'N/A'}")
        
        if not controller_found:
            print("  未找到ArticulationController节点")
        
    except Exception as e:
        print(f"获取Action Graph信息时出错: {e}")
        
    print("\n========================================")


# 使用视觉脚本检查Action Graph
def print_action_graph_visual_info():
    """通过NVIDIA Omniverse视觉脚本API获取Action Graph信息"""
    print("\n========== Action Graph信息 (视觉脚本API) ==========")
    
    try:
        # 尝试使用视觉脚本API
        try:
            import omni.kit.visual_scripting as vs
            print("成功导入视觉脚本API")
            
            # 尝试获取视觉脚本图表
            graphs = vs.get_all_graphs()
            if graphs:
                print(f"找到 {len(graphs)} 个视觉脚本图表")
                for i, graph in enumerate(graphs):
                    print(f"\n图表 {i+1}:")
                    print(f"  路径: {graph.get_path()}")
                    
                    # 获取节点
                    nodes = graph.get_nodes()
                    print(f"  节点数量: {len(nodes)}")
                    for node in nodes:
                        print(f"    - {node.get_title()} ({node.get_type_name()})")
                        
                        # 获取节点参数
                        params = node.get_parameters()
                        if params:
                            print("      参数:")
                            for param_name, param_value in params.items():
                                print(f"        {param_name}: {param_value}")
            else:
                print("没有找到视觉脚本图表")
                
        except ImportError:
            print("视觉脚本API不可用")
            
        # 尝试使用新版的图表API
        try:
            import omni.graph.core as og
            print("\n尝试使用omni.graph.core API")
            
            # 在Isaac Sim 4.2中，API可能发生了变化
            # 尝试使用可能的替代方法
            if hasattr(og, "get_graph_keys"):
                keys = og.get_graph_keys()
                print(f"找到 {len(keys)} 个图表键")
                
                for key in keys:
                    print(f"  图表键: {key}")
                    graph = og.get_graph(key)
                    if graph:
                        print("  成功获取图表")
                        # 处理图表...
            else:
                print("  get_graph_keys 方法不可用")
                
            # 尝试其他可能的方法
            if hasattr(og, "get_registered_graphs"):
                graphs = og.get_registered_graphs()
                print(f"找到 {len(graphs)} 个注册的图表")
                # 处理图表...
            else:
                print("  get_registered_graphs 方法不可用")
                
        except Exception as e:
            print(f"使用omni.graph.core时出错: {e}")
            
    except Exception as e:
        print(f"获取视觉脚本信息时出错: {e}")
        
    print("\n========================================")


# 使用USD直接查询机器人关节
def print_robot_joints_via_usd(robot_prim_path="/World/piper_description"):
    """使用USD API直接查询机器人关节信息"""
    print("\n========== USD API查询机器人关节信息 ==========")
    
    try:
        from pxr import Usd, UsdPhysics, UsdGeom, Gf
        
        # 获取当前舞台
        stage = omni.usd.get_context().get_stage()
        
        # 获取机器人根节点
        robot_prim = stage.GetPrimAtPath(robot_prim_path)
        if not robot_prim:
            print(f"找不到机器人: {robot_prim_path}")
            return
            
        print(f"找到机器人: {robot_prim_path}")
        
        # 查找关节
        joints = []
        
        # 方法1: 直接查找关节类型
        for prim in Usd.PrimRange(robot_prim):
            if prim.IsA(UsdPhysics.Joint) or prim.IsA(UsdPhysics.RevoluteJoint) or prim.IsA(UsdPhysics.PrismaticJoint):
                joints.append(prim)
        
        print(f"\n找到 {len(joints)} 个关节")
        
        if joints:
            for i, joint in enumerate(joints):
                joint_path = joint.GetPath()
                joint_name = joint.GetName()
                
                print(f"\n关节 {i+1}:")
                print(f"  路径: {joint_path}")
                print(f"  名称: {joint_name}")
                
                # 获取关节类型
                joint_type = "未知"
                if joint.IsA(UsdPhysics.RevoluteJoint):
                    joint_type = "旋转关节"
                elif joint.IsA(UsdPhysics.PrismaticJoint):
                    joint_type = "平移关节"
                elif joint.IsA(UsdPhysics.Joint):
                    joint_type = "通用关节"
                
                print(f"  类型: {joint_type}")
                
                # 获取关节限制
                if joint.IsA(UsdPhysics.RevoluteJoint):
                    joint_api = UsdPhysics.RevoluteJoint(joint)
                    lower = joint_api.GetLowerAttr().Get() if joint_api.GetLowerAttr().HasValue() else float('-inf')
                    upper = joint_api.GetUpperAttr().Get() if joint_api.GetUpperAttr().HasValue() else float('inf')
                    print(f"  限制范围: [{lower:.4f}, {upper:.4f}]")
                
                # 获取关节属性
                attrs = joint.GetAttributes()
                if attrs:
                    print("  属性:")
                    for attr in attrs:
                        if attr.HasValue() and attr.IsValid():
                            print(f"    - {attr.GetName()}: {attr.Get()}")
                
        else:
            print("未找到任何关节，尝试查找所有可能的关节...")
            
            # 方法2: 查找包含'joint'的节点
            possible_joints = []
            
            for prim in Usd.PrimRange(robot_prim):
                if "joint" in prim.GetName().lower():
                    possible_joints.append(prim)
            
            print(f"\n找到 {len(possible_joints)} 个可能的关节")
            
            for i, joint in enumerate(possible_joints[:5]):  # 只显示前5个
                print(f"  {i+1}. {joint.GetPath()}")
                
            if len(possible_joints) > 5:
                print(f"  ...以及 {len(possible_joints) - 5} 个其他关节")
                
    except Exception as e:
        print(f"使用USD API查询关节信息时出错: {e}")
        
    print("\n========================================")


# 综合打印物理模拟状态
def print_physics_status():
    """打印物理模拟状态"""
    print("\n========== 物理模拟状态检查 ==========")
    
    # 检查物理模拟状态
    timeline = omni.timeline.get_timeline_interface()
    is_playing = timeline.is_playing()
    print(f"物理模拟状态: {'运行中' if is_playing else '已停止'}")
    
    # 获取PhysX版本信息
    try:
        import omni.physx
        print("成功导入omni.physx")
        
        # 尝试获取PhysX接口
        physx_interface = omni.physx.get_physx_interface()
        if physx_interface:
            print("成功获取PhysX接口")
            
            # 尝试获取PhysX属性
            try:
                physx_props = {}
                
                # 测试可能的方法
                if hasattr(physx_interface, "get_scene_count"):
                    physx_props["场景数量"] = physx_interface.get_scene_count()
                    
                if hasattr(physx_interface, "is_simulating"):
                    physx_props["正在模拟"] = physx_interface.is_simulating()
                    
                if hasattr(physx_interface, "get_active_stage_id"):
                    physx_props["活动场景ID"] = physx_interface.get_active_stage_id()
                    
                if hasattr(physx_interface, "get_gravitational_acceleration"):
                    physx_props["重力加速度"] = physx_interface.get_gravitational_acceleration()
                
                if physx_props:
                    print("\nPhysX属性:")
                    for prop, value in physx_props.items():
                        print(f"  {prop}: {value}")
                else:
                    print("未找到可查询的PhysX属性")
            except Exception as e:
                print(f"获取PhysX属性时出错: {e}")
                
    except Exception as e:
        print(f"获取PhysX接口时出错: {e}")
    
    print("\n通过USD API查询物理属性:")
    try:
        from pxr import Usd, UsdPhysics
        
        # 获取当前舞台
        stage = omni.usd.get_context().get_stage()
        
        # 查找物理场景
        scene = UsdPhysics.Scene.Get(stage)
        if scene:
            scene_prim = scene.GetPrim()
            print(f"找到物理场景: {scene_prim.GetPath()}")
            
            # 获取重力
            gravity_attr = scene.GetGravityDirectionAttr()
            gravity_mag_attr = scene.GetGravityMagnitudeAttr()
            
            if gravity_attr and gravity_attr.HasValue():
                gravity_dir = gravity_attr.Get()
                print(f"  重力方向: {gravity_dir}")
                
            if gravity_mag_attr and gravity_mag_attr.HasValue():
                gravity_mag = gravity_mag_attr.Get()
                print(f"  重力大小: {gravity_mag}")
        else:
            print("未找到物理场景")
            
    except Exception as e:
        print(f"查询USD物理属性时出错: {e}")
        
    print("========================================")


# 执行所有函数
try:
    # 确保物理模拟已启动
    timeline = omni.timeline.get_timeline_interface()
    if not timeline.is_playing():
        timeline.play()
        print("已启动物理模拟")
    
    # 打印机器人信息
    print_robot_actions_info("/World/piper_description")
    
    # 使用兼容方式打印Action Graph信息
    print_action_graph_info_compatible()
    
    # 使用视觉脚本API获取信息
    print_action_graph_visual_info()
    
    # 使用USD API查询关节信息
    print_robot_joints_via_usd("/World/piper_description")
    
    # 打印物理模拟状态
    print_physics_status()
    
except Exception as e:
    print(f"执行过程中出现错误: {e}")