# probe_pose.py — 一次性探针：打印 cube 复位后的真实位姿 + 确认 success 张量
import argparse, sys
from isaaclab.app import AppLauncher
import cli_args  # isort: skip

parser = argparse.ArgumentParser()
parser.add_argument("--num_envs", type=int, default=4)
parser.add_argument("--task", type=str, default=None)
parser.add_argument("--agent", type=str, default="rsl_rl_cfg_entry_point")
parser.add_argument("--settle_steps", type=int, default=30, help="reset 后空走几步让物体落定")
cli_args.add_rsl_rl_args(parser)
AppLauncher.add_app_launcher_args(parser)
args_cli, hydra_args = parser.parse_known_args()
sys.argv = [sys.argv[0]] + hydra_args

app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

import gymnasium as gym
import torch
from isaaclab.envs import ManagerBasedRLEnvCfg
import isaaclab_tasks  # noqa: F401
import uwlab_tasks     # noqa: F401
from uwlab_tasks.utils.hydra import hydra_task_config


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg, agent_cfg):
    env_cfg.scene.num_envs = args_cli.num_envs
    env = gym.make(args_cli.task, cfg=env_cfg)
    u = env.unwrapped

    # 动作维度（用零动作空走，让物体在物理下落定）
    act_dim = u.action_manager.total_action_dim
    u.reset()
    zeros = torch.zeros((u.num_envs, act_dim), device=u.device)
    for _ in range(args_cli.settle_steps):
        u.step(zeros)

    origins = u.scene.env_origins                      # [N,3] 每个 env 的世界原点
    ins = u.scene["insertive_object"].data             # 要叠的 cube
    rec = u.scene["receptive_object"].data             # 目标 cube

    torch.set_printoptions(precision=4, sci_mode=False)
    print("\n==================== POSE PROBE ====================")
    print("env_origins (world):\n", origins)
    print("\ninsertive_object root_pos_w (world):\n", ins.root_pos_w)
    print("insertive_object pos REL env_origin:\n", ins.root_pos_w - origins)
    print("insertive_object root_quat_w (wxyz):\n", ins.root_quat_w)
    print("\nreceptive_object root_pos_w (world):\n", rec.root_pos_w)
    print("receptive_object pos REL env_origin:\n", rec.root_pos_w - origins)
    print("receptive_object root_quat_w (wxyz):\n", rec.root_quat_w)

    # 确认 success 张量可取
    pc = u.reward_manager.get_term_cfg("progress_context").func
    print("\nsuccess tensor:", pc.success.shape, pc.success.dtype, "->", pc.success)
    print("continuous_success_counter ->", pc.continuous_success_counter)
    print("====================================================\n")
    env.close()


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