from isaaclab.app import AppLauncher
import argparse

# -----------------------------------------------------------------------------
# 1. 启动 Isaac Sim
# -----------------------------------------------------------------------------
parser = argparse.ArgumentParser()
AppLauncher.add_app_launcher_args(parser)
args = parser.parse_args()

app_launcher = AppLauncher(args)
simulation_app = app_launcher.app

# -----------------------------------------------------------------------------
# 2. 必须先创建 SimulationContext（PhysX 在这里初始化）
# -----------------------------------------------------------------------------
import isaaclab.sim as sim_utils

sim_cfg = sim_utils.SimulationCfg(
    dt=1 / 120,
    device=args.device,
)
sim = sim_utils.SimulationContext(sim_cfg)

# （可选）设置相机
sim.set_camera_view(
    eye=[3.0, 0.0, 2.0],
    target=[0.0, 0.0, 0.5],
)

# -----------------------------------------------------------------------------
# 3. 再 play timeline（这一步非常关键）
# -----------------------------------------------------------------------------
import omni.timeline

timeline = omni.timeline.get_timeline_interface()
timeline.play()

# -----------------------------------------------------------------------------
# 4. 现在才可以安全创建 Env
# -----------------------------------------------------------------------------
import torch
from isaaclab.envs import ManagerBasedRLEnv
from mindbot.tasks.manager_based.pull.pull_env_cfg import PullEnvCfg


def main():
    env_cfg = PullEnvCfg()
    env_cfg.scene.num_envs = 1

    env = ManagerBasedRLEnv(cfg=env_cfg)

    obs, _ = env.reset()
    print("[OK] Env reset")

    for step in range(100000):
        actions = torch.zeros(
            env.num_envs,
            env.action_manager.total_action_dim,
            device=env.device,
        )
        obs, rew, done, trunc, info = env.step(actions)

        if step % 100 == 0:
            root_pos = env.scene["Mindbot"].data.root_pos_w[0]
            print(f"step {step}, root z = {root_pos[2].item():.3f}")

    print("[DONE]")


if __name__ == "__main__":
    main()
    simulation_app.close()