"""Probe the YAM bimanual env: load, reset, report action space + robot/object poses."""
import argparse, sys, os
from isaaclab.app import AppLauncher
parser = argparse.ArgumentParser()
AppLauncher.add_app_launcher_args(parser)
args = parser.parse_args()
args.headless = True; args.enable_cameras = True
app = AppLauncher(args).app

import numpy as np, torch, gymnasium as gym
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(REPO, "source"))
import bimanual.tasks.manager_based.yam  # noqa: F401  (registers Template-YAM-Play-v0)
from isaaclab_tasks.utils import parse_env_cfg

TASK = "Template-YAM-Play-v0"
cfg = parse_env_cfg(TASK, device="cuda:0", num_envs=1)
print("[yam] making env ...", flush=True)
env = gym.make(TASK, cfg=cfg)
u = env.unwrapped
obs, _ = env.reset()
print("[yam] action_space:", env.action_space, flush=True)
print("[yam] scene keys:", list(u.scene.keys()), flush=True)
origin = u.scene.env_origins[0].detach().cpu().numpy()
for name in ["left_robot", "right_robot"]:
    if name in u.scene.articulations:
        a = u.scene.articulations[name]
        p = a.data.root_pos_w[0].detach().cpu().numpy() - origin
        bns = list(a.data.body_names)
        print(f"[yam] {name} base(env-local)={np.round(p,3)} bodies={bns}", flush=True)
        print(f"[yam]   joints={list(a.data.joint_names)}", flush=True)
for name in ["apple", "grape", "cart"]:
    if name in u.scene.rigid_objects:
        p = u.scene.rigid_objects[name].data.root_pos_w[0].detach().cpu().numpy() - origin
        print(f"[yam] {name}(env-local)={np.round(p,3)}", flush=True)
# step a couple zero actions
adim = env.action_space.shape[-1]
print("[yam] action_dim:", adim, flush=True)
env.close(); app.close(); print("YAM_PROBE_OK", flush=True)