"""Print the world-bbox size of scene objects, so grasp widths and scales can be set from the
real geometry instead of by feel.

    python scripts/yam_size_probe.py --headless --objects grape,apple,can \
        --kit_args="--/rtx/verifyDriverVersion/enabled=false"
"""
import argparse, sys, os
from isaaclab.app import AppLauncher
p = argparse.ArgumentParser()
p.add_argument("--objects", default="grape,apple,can,peg")
AppLauncher.add_app_launcher_args(p)
args = p.parse_args(); args.headless = True
app = AppLauncher(args).app
import numpy as np, gymnasium as gym
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(REPO, "source")); sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import bimanual.tasks.manager_based.yam  # noqa
from isaaclab_tasks.utils import parse_env_cfg

TASK = "Template-YAM-Play-v0"; dev = "cuda:0"
cfg = parse_env_cfg(TASK, device=dev, num_envs=1)
env = gym.make(TASK, cfg=cfg, render_mode=None); u = env.unwrapped; env.reset()
import omni.usd
from pxr import UsdGeom, Usd
stage = omni.usd.get_context().get_stage()
bb = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
JAW_MAX = 2*0.04695   # both prismatic fingers fully open
print(f"[size] jaw max opening = {JAW_MAX:.4f} m", flush=True)
for name in [n for n in args.objects.split(",") if n]:
    if name not in u.scene.rigid_objects:
        print(f"[size] {name}: NOT IN SCENE", flush=True); continue
    pp = u.scene.rigid_objects[name].root_physx_view.prim_paths[0]
    rng = bb.ComputeWorldBound(stage.GetPrimAtPath(pp)).ComputeAlignedRange()
    ext = np.array(rng.GetMax())-np.array(rng.GetMin())
    # the top-down grasp closes along world y, so that extent is what must fit the jaw
    print(f"[size] {name:10s} extents(x,y,z)={np.round(ext,4)}  grasp_width(y)={ext[1]:.4f} "
          f"({100*ext[1]/JAW_MAX:.0f}% of jaw)", flush=True)
env.close(); app.close(); print("SIZE_PROBE_OK", flush=True)