from omni.isaac.kit import SimulationApp
# 启动 headless，确保和批量捕获脚本相同的执行环境
simulation_app = SimulationApp({"headless": True})

import omni.usd
from pxr import UsdGeom
import numpy as np
from scipy.spatial.transform import Rotation as R
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--scene_usd", required=True, help="USD 场景文件路径")
    parser.add_argument("--prim_path", default="/World/MySensorCamera", help="相机 prim 路径")
    args = parser.parse_args()

    usd_path = args.scene_usd
    prim_path = args.prim_path

    # 打开 USD 场景
    omni.usd.get_context().close_stage()
    success = omni.usd.get_context().open_stage(usd_path)
    if not success:
        print(f"❌ 无法打开 {usd_path}")
        simulation_app.close()
        return

    # 获取 stage
    stage = omni.usd.get_context().get_stage()
    prim = stage.GetPrimAtPath(prim_path)

    if not prim.IsValid():
        print(f"❌ 找不到 prim: {prim_path}")
        simulation_app.close()
        return

    # 读取世界变换矩阵
    xform = UsdGeom.Xformable(prim)
    world_xform = xform.ComputeLocalToWorldTransform(0)  # 0 代表默认时间码

    # 提取位置
    pos = [world_xform.ExtractTranslation()[i] for i in range(3)]

    # 提取四元数 [x, y, z, w] (USD标准: imag=(x,y,z), real=w)
    rot = world_xform.ExtractRotationQuat()
    quat = np.array([
        rot.GetImaginary()[0],
        rot.GetImaginary()[1],
        rot.GetImaginary()[2],
        rot.GetReal()
    ])

    # 打印结果
    print("==== Headless 模式读取相机 ====")
    print(f"Scene: {usd_path}")
    print(f"Prim: {prim_path}")
    print(f"World Position: {pos}")
    print(f"World Quaternion [x,y,z,w]: {quat.tolist()}")
    print(f"Euler angles (xyz deg): {R.from_quat(quat).as_euler('xyz', degrees=True).tolist()}")

    simulation_app.close()

if __name__ == "__main__":
    main()