"""Find the rotation that makes a RoboTwin asset sit the right way up.

RoboTwin GLBs are not authored Z-up, so a cup spawns upside-down and a basket spawns on its
side. Rather than guess one rotation per asset, this drops the object at several candidate
orientations, lets each settle, and RENDERS a labelled contact-sheet so the correct one can be
picked by eye -- plus the settled bbox height, which is the objective tell (an upright cup is
taller than it is wide; a tipped one is not).

    ROBOTWIN_USD=... YAM_RW_OBJECTS="rw_cupA=021_cup/base0.usd:0.075:0.08" \
    python scripts/yam_orient_probe.py --headless --obj rw_cupA --out outputs/probe_cup.png
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

p = argparse.ArgumentParser()
p.add_argument("--obj", required=True)
p.add_argument("--xy", default="0.00,0.10")
p.add_argument("--rpys", default="0,0,0;180,0,0;90,0,0;-90,0,0;0,90,0;0,180,0",
               help="';'-separated roll,pitch,yaw candidates in degrees")
p.add_argument("--out", default="outputs/orient_probe.png")
AppLauncher.add_app_launcher_args(p)
args = p.parse_args(); args.headless = True; args.enable_cameras = True
app = AppLauncher(args).app

import numpy as np, torch, gymnasium as gym
from PIL import Image, ImageDraw
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)
cfg.episode_length_s = 1.0e6
try:
    cfg.terminations.time_out = None
except Exception:
    pass
try:
    # close-in view of the object spot, so the contact sheet actually shows the orientation
    cfg.viewer.eye = (0.55, -0.35, 0.85); cfg.viewer.lookat = (0.0, 0.10, 0.50)
    cfg.viewer.resolution = (640, 480)
except Exception as e:
    print("viewer cfg:", e)
env = gym.make(TASK, cfg=cfg, render_mode="rgb_array"); u = env.unwrapped; env.reset()
if args.obj not in u.scene.rigid_objects:
    raise SystemExit(f"[probe] {args.obj!r} not in scene; register it with YAM_RW_OBJECTS")
OBJ = u.scene.rigid_objects[args.obj]
origin = u.scene.env_origins[0].cpu().numpy()
TABLE_TOP = 0.45
XY = [float(v) for v in args.xy.split(",")]

import omni.usd
from pxr import UsdGeom, Usd
stage = omni.usd.get_context().get_stage()
bbc = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
rng = bbc.ComputeWorldBound(stage.GetPrimAtPath(OBJ.root_physx_view.prim_paths[0])).ComputeAlignedRange()
ext = np.array(rng.GetMax())-np.array(rng.GetMin())
print(f"[probe] {args.obj} authored extents={np.round(ext,4)}", flush=True)


def quat(r, pi, y):
    r, pi, y = np.radians([r, pi, y])
    cr, sr = np.cos(r/2), np.sin(r/2); cp, sp = np.cos(pi/2), np.sin(pi/2); cy, sy = np.cos(y/2), np.sin(y/2)
    return np.array([cr*cp*cy+sr*sp*sy, sr*cp*cy-cr*sp*sy, cr*sp*cy+sr*cp*sy, cr*cp*sy-sr*sp*cy])


def zeros_act():
    return torch.zeros((1, env.action_space.shape[1]), dtype=torch.float32, device=dev)


tiles = []
for spec in [s for s in args.rpys.split(";") if s]:
    r, pi, y = [float(v) for v in spec.split(",")]
    q = quat(r, pi, y)
    seat = TABLE_TOP+float(max(ext))/2.0+0.01          # drop from just above, whatever axis is up
    OBJ.write_root_pose_to_sim(torch.tensor(np.concatenate([origin+np.array([XY[0], XY[1], seat]), q]),
                                            dtype=torch.float32, device=dev).view(1, 7))
    OBJ.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))
    for _ in range(120):                                # let it settle onto the table
        env.step(zeros_act())
    w = OBJ.data.root_pos_w[0].cpu().numpy()-origin
    qs = OBJ.data.root_quat_w[0].cpu().numpy()
    # world-space height of the object above the table = the objective "is it upright" signal
    height = float(w[2]-TABLE_TOP)*2
    img = env.render()
    im = Image.fromarray(np.asarray(img)[..., :3].copy()) if img is not None else Image.new("RGB", (640, 480))
    d = ImageDraw.Draw(im)
    txt = f"rpy=({r:.0f},{pi:.0f},{y:.0f})  settled z={w[2]:.3f} (~{height:.3f} tall)"
    d.rectangle([0, 0, 430, 22], fill=(0, 0, 0)); d.text((5, 4), txt, fill=(255, 235, 60))
    tiles.append(im)
    print(f"[probe] rpy=({r:.0f},{pi:.0f},{y:.0f}) settled pos=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f}) "
          f"quat={np.round(qs,3)}", flush=True)

cols = 3
rows = (len(tiles)+cols-1)//cols
W, H = tiles[0].size
sheet = Image.new("RGB", (W*cols, H*rows))
for i, im in enumerate(tiles):
    sheet.paste(im, (W*(i % cols), H*(i//cols)))
os.makedirs(os.path.dirname(args.out), exist_ok=True)
sheet.save(args.out)
print(f"[probe] contact sheet -> {args.out}", flush=True)
env.close(); app.close(); print("ORIENT_PROBE_OK", flush=True)