# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Script to run an environment with zero action agent."""

"""Launch Isaac Sim Simulator first."""

import argparse

from isaaclab.app import AppLauncher

# add argparse arguments
parser = argparse.ArgumentParser(description="Zero agent for Isaac Lab environments.")
parser.add_argument(
    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
)
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()

# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

"""Rest everything follows."""

import gymnasium as gym
import torch

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import parse_env_cfg

import mindbot.tasks  # noqa: F401


def main():
    """Zero actions agent with Isaac Lab environment."""
    # parse configuration
    env_cfg = parse_env_cfg(
        args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric
    )
    # create environment
    env = gym.make(args_cli.task, cfg=env_cfg)

    # print info (this is vectorized environment)
    print(f"[INFO]: Gym observation space: {env.observation_space}")
    print(f"[INFO]: Gym action space: {env.action_space}")
    # reset environment
    env.reset()
    
    # 计数器：每 120 步打印一次
    step_count = 0
    print_interval = 120
    
    # simulate environment
    while simulation_app.is_running():
        # run everything in inference mode
        with torch.inference_mode():
            # compute zero actions
            actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device)
            # apply actions
            env.step(actions)
            
            # ... (前面的代码保持不变)

            # 每 120 步打印一次坐标
            step_count += 1
            if step_count % print_interval == 0:
                scene = env.unwrapped.scene
                
                # 1. 获取 centrifuge 的世界坐标
                try:
                    centrifuge = scene["centrifuge"]
                    centrifuge_pos = centrifuge.data.root_pos_w[0].cpu().numpy()
                    centrifuge_quat = centrifuge.data.root_quat_w[0].cpu().numpy()
                    print(f"[Step {step_count}] CENTRIFUGE - Pos: {centrifuge_pos}, Quat: {centrifuge_quat}")
                except KeyError:
                    print(f"[Step {step_count}] CENTRIFUGE - Not found")

                # 2. 获取 Reservoir_A 的世界坐标
                # 注意：这里的 key ("reservoir_a") 必须与你在 SceneCfg 中定义的名称一致
                try:
                    # 如果你在配置里命名为 "reservoir_a"
                    reservoir = scene["reservoir_a"] 
                    res_pos = reservoir.data.root_pos_w[0].cpu().numpy()
                    res_quat = reservoir.data.root_quat_w[0].cpu().numpy()
                    print(f"[Step {step_count}] RESERVOIR_A - Pos: {res_pos}, Quat: {res_quat}")
                except KeyError:
                    # 如果 Reservoir_A 是 centrifuge 机器人(Articulation)的一个 Link（身体部件）
                    # 我们可以从 centrifuge 的 body 数据中获取
                    try:
                        centrifuge = scene["centrifuge"]
                        # 找到名为 'Reservoir_A' 的 link 索引
                        body_names = centrifuge.body_names
                        if "Reservoir_A" in body_names:
                            idx = body_names.index("Reservoir_A")
                            res_pos = centrifuge.data.body_pos_w[0, idx].cpu().numpy()
                            res_quat = centrifuge.data.body_quat_w[0, idx].cpu().numpy()
                            print(f"[Step {step_count}] RESERVOIR_A (Link) - Pos: {res_pos}, Quat: {res_quat}")
                        else:
                            print(f"[Step {step_count}] RESERVOIR_A - Not found in scene keys or links")
                    except Exception:
                        print(f"[Step {step_count}] RESERVOIR_A - Not found")
                
                # 3. 获取 lid 的世界坐标
                try:
                    lid = scene["lid"]
                    lid_pos = lid.data.root_pos_w[0].cpu().numpy()
                    lid_quat = lid.data.root_quat_w[0].cpu().numpy()
                    print(f"[Step {step_count}] LID - Pos: {lid_pos}, Quat: {lid_quat}")
                except KeyError:
                    print(f"[Step {step_count}] LID - Not found")
                
                print("-" * 80)

# ... (后面的代码保持不变)
            # 每 120 步打印一次坐标
            # step_count += 1
            # if step_count % print_interval == 0:
            #     scene = env.unwrapped.scene
                
            #     # 获取 centrifuge 的世界坐标（root position）
            #     try:
            #         centrifuge = scene["centrifuge"]
            #         centrifuge_pos = centrifuge.data.root_pos_w[0].cpu().numpy()  # 取第一个环境
            #         centrifuge_quat = centrifuge.data.root_quat_w[0].cpu().numpy()
            #         print(f"[Step {step_count}] CENTRIFUGE_CFG - World Position: {centrifuge_pos}, Quaternion: {centrifuge_quat}")
            #     except KeyError:
            #         print(f"[Step {step_count}] CENTRIFUGE_CFG - Not found in scene")
                
            #     # 获取 lid 的世界坐标
            #     try:
            #         lid = scene["lid"]
            #         lid_pos = lid.data.root_pos_w[0].cpu().numpy()  # 取第一个环境
            #         lid_quat = lid.data.root_quat_w[0].cpu().numpy()
            #         print(f"[Step {step_count}] LID_CFG - World Position: {lid_pos}, Quaternion: {lid_quat}")
            #     except KeyError:
            #         print(f"[Step {step_count}] LID_CFG - Not found in scene")
                
            #     print("-" * 80)

    # close the simulator
    env.close()


if __name__ == "__main__":
    # run the main function
    main()
    # close sim app
    simulation_app.close()

# # Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# # All rights reserved.
# #
# # SPDX-License-Identifier: BSD-3-Clause

# """Script to run an environment with zero action agent."""

# """Launch Isaac Sim Simulator first."""

# import argparse

# from isaaclab.app import AppLauncher

# # add argparse arguments
# parser = argparse.ArgumentParser(description="Zero agent for Isaac Lab environments.")
# parser.add_argument(
#     "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
# )
# parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
# parser.add_argument("--task", type=str, default=None, help="Name of the task.")
# # append AppLauncher cli args
# AppLauncher.add_app_launcher_args(parser)
# # parse the arguments
# args_cli = parser.parse_args()

# # launch omniverse app
# app_launcher = AppLauncher(args_cli)
# simulation_app = app_launcher.app

# """Rest everything follows."""

# import gymnasium as gym
# import torch

# import isaaclab_tasks  # noqa: F401
# from isaaclab_tasks.utils import parse_env_cfg

# import mindbot.tasks  # noqa: F401


# def main():
#     """Zero actions agent with Isaac Lab environment."""
#     # parse configuration
#     env_cfg = parse_env_cfg(
#         args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric
#     )
#     # create environment
#     env = gym.make(args_cli.task, cfg=env_cfg)

#     # print info (this is vectorized environment)
#     print(f"[INFO]: Gym observation space: {env.observation_space}")
#     print(f"[INFO]: Gym action space: {env.action_space}")
#     # reset environment
#     env.reset()
#     # simulate environment

#     step_count = 0
#     print_interval = 120

#     while simulation_app.is_running():
#         # run everything in inference mode
#         with torch.inference_mode():
#             # compute zero actions
#             actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device)
#             # apply actions
#             env.step(actions)

#     # close the simulator
#     env.close()


# if __name__ == "__main__":
#     # run the main function
#     main()
#     # close sim app
#     simulation_app.close()
