"""YAM BIMANUAL pot lift: both arms grasp the pot by its two opposite handles and lift it
together, keeping it level.

One arm cannot lift a wide pot -- the point of the task is the two-arm synchronisation: both
jaws close on opposite handles, then both wrists rise along the same profile so the pot stays
flat. The pot is a RoboTwin (MIT) asset converted GLB->USD; its handle axis is measured from
the mesh bbox at runtime and the pot is yawed so the handles face the two arms.

    ROBOTWIN_USD=/path/to/usd YAM_RW_OBJECTS="rw_pot=060_kitchenpot/base0.usd:0.05:0.30" \
    python scripts/yam_pot_lift.py --headless --video outputs/tasks/pot.mp4
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser()
parser.add_argument("--obj", default="rw_pot", help="pot object registered via YAM_RW_OBJECTS")
parser.add_argument("--pot_xy", default="0.04,0.00", help="pot centre x,y (env-local)")
parser.add_argument("--pot_rpy", default="0,0,0",
                    help="initial roll,pitch,yaw in DEGREES; RoboTwin meshes are not all authored "
                         "z-up, so a pot may need a roll to stand on its base")
parser.add_argument("--grip_tilt", type=float, default=0.0,
                    help="tilt each gripper this many DEGREES away from straight-down, leaning in "
                         "over the pot, so the fingers hook around an ear handle from outside "
                         "instead of pressing onto it from above")
parser.add_argument("--half_y", type=float, default=-1.0,
                    help="override the handle half-span (m). The USD bbox is the AUTHORED one and "
                         "does not follow the mesh once it is rotated, so measuring it is unreliable")
parser.add_argument("--grip_z", type=float, default=-1.0, help="override the grip height (world z)")
parser.add_argument("--grip_max", type=float, default=0.022,
                    help="a stall wider than this (per finger) means the jaw hit the body, not a "
                         "handle; treat it as no grasp rather than a successful clamp")
parser.add_argument("--half_y_l", type=float, default=-1.0,
                    help="per-arm handle offset for the LEFT arm; the mesh origin is not centred "
                         "between the two handles, so a single symmetric span misses one of them")
parser.add_argument("--half_y_r", type=float, default=-1.0, help="per-arm handle offset, RIGHT arm")
parser.add_argument("--jaw", default="y", choices=["x", "y"],
                    help="world axis each jaw closes along. 'y' grips a pot rim/handle from the "
                         "side; 'x' grips ACROSS a plank's width, which is how two arms carry a board")
parser.add_argument("--probe_only", action="store_true",
                    help="place the pot, let it settle, print its resting size/orientation, exit")
parser.add_argument("--lift", type=float, default=0.14, help="how high to lift the pot (m)")
parser.add_argument("--grip_inset", type=float, default=0.004,
                    help="close the jaws this far inside the handle tip (m)")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/yam_pot_lift.mp4")
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
import imageio.v2 as imageio
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 yam_prm  # noqa: F401  (kept for parity with the other demos)
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              # scripted sequence: never auto-reset mid-run
try:
    _cfg.terminations.time_out = None
except Exception as _e:
    print("[pot] time_out disable failed:", _e)
try:
    _cfg.viewer.eye = (0.95, -0.95, 1.15); _cfg.viewer.lookat = (0.05, 0.0, 0.5)
    _cfg.viewer.resolution = (720, 540)
except Exception as _e:
    print("viewer cfg:", _e)
env = gym.make(TASK, cfg=_cfg, render_mode="rgb_array"); u = env.unwrapped; env.reset()


def Rq(q):
    w, x, y, z = q
    return np.array([[1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w)],
                     [2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w)],
                     [2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)]])


def qR(m):
    t = m[0, 0]+m[1, 1]+m[2, 2]
    if t > 0:
        s = np.sqrt(t+1)*2; w = .25*s; x = (m[2, 1]-m[1, 2])/s; y = (m[0, 2]-m[2, 0])/s; z = (m[1, 0]-m[0, 1])/s
    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
        s = np.sqrt(1+m[0, 0]-m[1, 1]-m[2, 2])*2; w = (m[2, 1]-m[1, 2])/s; x = .25*s; y = (m[0, 1]+m[1, 0])/s; z = (m[0, 2]+m[2, 0])/s
    elif m[1, 1] > m[2, 2]:
        s = np.sqrt(1+m[1, 1]-m[0, 0]-m[2, 2])*2; w = (m[0, 2]-m[2, 0])/s; x = (m[0, 1]+m[1, 0])/s; y = .25*s; z = (m[1, 2]+m[2, 1])/s
    else:
        s = np.sqrt(1+m[2, 2]-m[0, 0]-m[1, 1])*2; w = (m[1, 0]-m[0, 1])/s; x = (m[0, 2]+m[2, 0])/s; y = (m[1, 2]+m[2, 1])/s; z = .25*s
    q = np.array([w, x, y, z]); q /= np.linalg.norm(q)+1e-9
    return q if q[0] >= 0 else -q


origin = u.scene.env_origins[0].cpu().numpy()
R = u.scene["right_robot"]; Rbn = list(R.data.body_names)
L = u.scene["left_robot"]; Lbn = list(L.data.body_names)
rroot = R.data.root_pos_w[0].cpu().numpy()-origin; rrootq = R.data.root_quat_w[0].cpu().numpy()
lroot = L.data.root_pos_w[0].cpu().numpy()-origin; lrootq = L.data.root_quat_w[0].cpu().numpy()
OFF = np.array([0, 0, 0.13]); TABLE_TOP = 0.45
OPEN, CLOSE = 1.0, -1.0
if args.obj not in u.scene.rigid_objects:
    raise SystemExit(f"[pot] object {args.obj!r} is not in the scene -- register it with "
                     f"YAM_RW_OBJECTS and point ROBOTWIN_USD at the converted assets")


def eef_root(a, bn, root, rootq):
    i = bn.index("link_6"); p = a.data.body_pos_w[0, i].cpu().numpy()-origin
    q = a.data.body_quat_w[0, i].cpu().numpy()
    return Rq(rootq).T@((p+Rq(q)@OFF)-root), q


lp0, lq0 = eef_root(L, Lbn, lroot, lrootq)
rp0, rq0 = eef_root(R, Rbn, rroot, rrootq)
def tilted_grasp(sign):
    """Grasp frame that leans `--grip_tilt` deg off vertical toward the pot's centre.

    `sign` is +1 for the arm on the +y side and -1 for the -y side, so both hands come at the
    ear handles from outside-in. Jaw closes along world x, i.e. ACROSS the ear, which is the
    only way a parallel jaw can hold a tab that sticks out sideways.
    """
    t = np.radians(args.grip_tilt)
    c, s = np.cos(t), np.sin(t)
    X = np.array([1., 0., 0.])                       # jaw axis: across the handle
    Z = np.array([0., -sign*s, -c])                  # approach: down, leaning toward the centre
    Y = np.cross(Z, X)
    return qR(np.stack([X, Y, Z], axis=1))


if args.grip_tilt != 0.0:
    gq = tilted_grasp(+1)                            # replaced per-arm below
else:
    gq = qR(np.stack([np.array([1., 0., 0.]), np.array([0., -1., 0.]), np.array([0., 0., -1.])], axis=1)
            if args.jaw == "x" else
            np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))
GQ_L = tilted_grasp(+1) if args.grip_tilt != 0.0 else gq     # left arm sits at +y
GQ_R = tilted_grasp(-1) if args.grip_tilt != 0.0 else gq     # right arm sits at -y


def act2(lp, lg, rp, rg):
    # each arm carries its own grasp orientation so the two hands can lean in symmetrically
    return torch.tensor(np.concatenate([lp, GQ_L, [lg], rp, GQ_R, [rg]]),
                        dtype=torch.float32, device=dev).view(1, -1)


# ---- place the pot, then measure it ------------------------------------------------------
PXY = [float(v) for v in args.pot_xy.split(",")]
pot = u.scene.rigid_objects[args.obj]


def _quat(roll, pitch, yaw):
    cr, sr = np.cos(roll/2), np.sin(roll/2)
    cp, sp = np.cos(pitch/2), np.sin(pitch/2)
    cy, sy = np.cos(yaw/2), np.sin(yaw/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])


_RPY = [np.radians(float(v)) for v in args.pot_rpy.split(",")]


def _set_pot(xy, yaw_extra, z=0.62):
    q = _quat(_RPY[0], _RPY[1], _RPY[2]+yaw_extra)
    pose = torch.tensor(np.concatenate([origin+np.array([xy[0], xy[1], z]), q]),
                        dtype=torch.float32, device=dev).view(1, 7)
    pot.write_root_pose_to_sim(pose); pot.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))


_set_pot(PXY, 0.0)
for _ in range(70):
    env.step(act2(lp0, OPEN, rp0, OPEN))

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(pot.root_physx_view.prim_paths[0])).ComputeAlignedRange()
ext = np.array(rng.GetMax())-np.array(rng.GetMin())
print(f"[pot] {args.obj} bbox extents={np.round(ext,3)}", flush=True)
# Handles stick out along the pot's LONG horizontal axis. The two arms work at -y (right) and
# +y (left), so if the long axis is x, yaw the pot 90 deg to face the handles at the arms.
# Re-seat it just above the table first: dropping a tall pot from spawn height topples it before
# the arms ever reach it. Only once it is settled is the bbox meaningful, so the handle-facing
# yaw is decided from THAT measurement (deciding from the mid-air bbox picked the wrong axis).
_set_pot(PXY, 0.0, z=TABLE_TOP+float(ext[2])/2.0+0.004)
for _ in range(90):
    env.step(act2(lp0, OPEN, rp0, OPEN))
rng = bbc.ComputeWorldBound(stage.GetPrimAtPath(pot.root_physx_view.prim_paths[0])).ComputeAlignedRange()
ext = np.array(rng.GetMax())-np.array(rng.GetMin())
# handles lie along the pot's LONG horizontal axis; the arms work at -y (right) and +y (left)
_yaw = np.pi/2 if ext[0] > ext[1]*1.05 else 0.0
if _yaw:
    _set_pot(PXY, _yaw, z=TABLE_TOP+float(ext[2])/2.0+0.004)
    for _ in range(90):
        env.step(act2(lp0, OPEN, rp0, OPEN))
    # BBoxCache reads the AUTHORED transform, not the physics pose, so re-measuring after the
    # yaw returns the same numbers. Swap the horizontal extents by hand instead -- otherwise the
    # arms aim 5.8 cm from centre (inside the pot body) rather than at the 12.6 cm handle tips.
    ext = np.array([ext[1], ext[0], ext[2]])
_q0 = pot.data.root_quat_w[0].cpu().numpy()
# "Upright" is measured against the pose it SETTLED in, not against world +z: a RoboTwin mesh
# needs a 90 deg roll to stand up, so its local +z is horizontal by construction and comparing
# to world +z reports a perfectly upright pot as "tilted 90 degrees".
_UP_REF = Rq(_q0)@np.array([0., 0., 1.])
_tilt0 = 0.0


def tilt_from_rest(q):
    v = Rq(q)@np.array([0., 0., 1.])
    return float(np.degrees(np.arccos(np.clip(float(np.dot(v, _UP_REF)), -1.0, 1.0))))
print(f"[pot] seated: yaw={np.degrees(_yaw):.0f}deg extents={np.round(ext,3)} rest_tilt={_tilt0:.1f}deg", flush=True)
if args.probe_only:
    _w = pot.data.root_pos_w[0].cpu().numpy()-origin
    print(f"[pot] PROBE pos=({_w[0]:.3f},{_w[1]:.3f},{_w[2]:.3f}) extents={np.round(ext,4)} "
          f"tilt={_tilt0:.1f}deg", flush=True)
    env.close(); app.close(); print("YAM_POT_PROBE_OK", flush=True); sys.exit(0)

pw = pot.data.root_pos_w[0].cpu().numpy()-origin
half_y = float(args.half_y) if args.half_y > 0 else float(ext[1])/2.0
pot_h = float(ext[2])
# Grip the handle tips: just inside the y extremes, at the pot's upper body where the handles are.
grip_z = float(args.grip_z) if args.grip_z > 0 else TABLE_TOP + pot_h*0.72
_hy_l = float(args.half_y_l) if args.half_y_l > 0 else half_y
_hy_r = float(args.half_y_r) if args.half_y_r > 0 else half_y
Ly = pw[1] + _hy_l - args.grip_inset      # left arm  -> +y handle
Ry = pw[1] - _hy_r + args.grip_inset      # right arm -> -y handle
L_grip_w = np.array([pw[0], Ly, grip_z], np.float32)
R_grip_w = np.array([pw[0], Ry, grip_z], np.float32)
print(f"[pot] pot at ({pw[0]:.3f},{pw[1]:.3f}) h={pot_h:.3f} half_y={half_y:.3f} grip_z={grip_z:.3f}", flush=True)


def to_root(w, root, rootq):
    return (Rq(rootq).T@(np.asarray(w, np.float32)-root)).astype(np.float32)


L_grip = to_root(L_grip_w, lroot, lrootq); R_grip = to_root(R_grip_w, rroot, rrootq)
L_pre = L_grip+np.array([0, 0, 0.13], np.float32); R_pre = R_grip+np.array([0, 0, 0.13], np.float32)
print(f"[pot] L_grip(root)={np.round(L_grip,3)}  R_grip(root)={np.round(R_grip,3)}", flush=True)

lhome_q = L.data.joint_pos[0].clone()


def _boost(view, tag, s=1.8, d=1.6):
    try:
        m = view.get_material_properties().clone(); m[..., 0] = s; m[..., 1] = d
        view.set_material_properties(m, torch.arange(m.shape[0], dtype=torch.int32, device=m.device))
    except Exception as e:
        print(f"[pot] friction set failed on {tag}:", e, flush=True)


_boost(R.root_physx_view, "right"); _boost(L.root_physx_view, "left")
_boost(pot.root_physx_view, args.obj)


def eefL():
    p, _ = eef_root(L, Lbn, lroot, lrootq); return p


def eefR():
    p, _ = eef_root(R, Rbn, rroot, rrootq); return p


def potw():
    return pot.data.root_pos_w[0].cpu().numpy()-origin


def potq():
    return pot.data.root_quat_w[0].cpu().numpy()


def fsep(art):
    jn = list(art.data.joint_names)
    return (float(art.data.joint_pos[0, jn.index("left_finger")].item())
            + float(art.data.joint_pos[0, jn.index("right_finger")].item()))/2


# ---- overlay + smooth two-arm executor ---------------------------------------------------
frames = []; _phase = {"v": "approach"}; _RESULT = {"v": ""}; _GRIP = {"v": "OPEN"}
SEM = {"approach": "1. BOTH ARMS APPROACH handles", "descend": "2. DESCEND to handle height",
       "close": "3. BOTH JAWS CLOSE on handles", "lift": "4. SYNCHRONISED LIFT",
       "hold": "5. HOLD (pot level)", "lower": "6. LOWER back", "done": "DONE"}


def capture():
    img = env.render()
    if img is None:
        return
    im = Image.fromarray(np.asarray(img)[..., :3].copy()); d = ImageDraw.Draw(im)
    w = potw()
    lines = ["=== DUAL-ARM POT LIFT (two handles) ==="
             + (f"  EP {args.episode}" if args.episode >= 0 else "")]
    if _RESULT["v"]:
        lines.append(f"RESULT: {_RESULT['v']}")
    lines += [f"ACTION: {SEM.get(_phase['v'], _phase['v'])}",
              f"grippers={_GRIP['v']}   pot_z={w[2]:.3f} (start {z0:.3f})" if "z0" in globals()
              else f"grippers={_GRIP['v']}   pot_z={w[2]:.3f}"]
    d.rectangle([0, 0, 420, 18*len(lines)+6], fill=(0, 0, 0))
    y = 3
    for ln in lines:
        d.text((6, y), ln, fill=(255, 235, 60)); y += 18
    frames.append(np.array(im))


_CORRL = {"v": np.zeros(3, np.float32)}; _CORRR = {"v": np.zeros(3, np.float32)}
_CMD = {"l": None, "r": None}


def _ease(a):
    return float(0.5-0.5*np.cos(np.pi*min(max(a, 0.0), 1.0)))


def drive(lt, lg, rt, rg, n, hold_left=False):
    """Both arms LERP to their targets on the SAME profile, so the pot stays level."""
    ls = _CMD["l"].copy() if _CMD["l"] is not None else eefL().astype(np.float32)
    rs = _CMD["r"].copy() if _CMD["r"] is not None else eefR().astype(np.float32)
    lt = np.asarray(lt, np.float32); rt = np.asarray(rt, np.float32)
    _GRIP["v"] = ("CLOSE" if lg < 0 else "OPEN")
    cl = _CORRL["v"]; cr = _CORRR["v"]
    for k in range(n):
        a = _ease((k+1)/float(n))
        lc = (1-a)*ls+a*lt; rc = (1-a)*rs+a*rt
        _CMD["l"] = lc; _CMD["r"] = rc
        env.step(act2((lc+cl).astype(np.float32), lg, (rc+cr).astype(np.float32), rg))
        if hold_left:
            L.write_joint_state_to_sim(lhome_q.view(1, -1), torch.zeros((1, lhome_q.shape[0]), device=dev))
        el = lc-eefL(); el = np.where(np.abs(el) > 0.008, el, 0.0)
        er = rc-eefR(); er = np.where(np.abs(er) > 0.008, er, 0.0)
        cl = np.clip(cl+0.08*el, -0.10, 0.10); cl[2] = max(float(cl[2]), -0.06)
        cr = np.clip(cr+0.08*er, -0.10, 0.10); cr[2] = max(float(cr[2]), -0.06)
        _CORRL["v"] = cl; _CORRR["v"] = cr
        if k % 3 == 0:
            capture()


z0 = float(potw()[2])
_phase["v"] = "approach"; drive(L_pre, OPEN, R_pre, OPEN, 110)
print(f"[pot] approached: L_err={np.linalg.norm(eefL()-L_pre):.3f} R_err={np.linalg.norm(eefR()-R_pre):.3f}", flush=True)
_phase["v"] = "descend"; drive(L_grip, OPEN, R_grip, OPEN, 120)
print(f"[pot] descended: L_err={np.linalg.norm(eefL()-L_grip):.3f} R_err={np.linalg.norm(eefR()-R_grip):.3f}", flush=True)

_phase["v"] = "close"; _GRIP["v"] = "CLOSE"
Lhold = eefL().astype(np.float32); Rhold = eefR().astype(np.float32)
_CMD["l"] = Lhold.copy(); _CMD["r"] = Rhold.copy()
prevL, prevR = fsep(L), fsep(R); stall = 0
for k in range(170):
    env.step(act2((Lhold+_CORRL["v"]).astype(np.float32), CLOSE, (Rhold+_CORRR["v"]).astype(np.float32), CLOSE))
    if k % 5 == 0:
        capture()
    curL, curR = fsep(L), fsep(R)
    stall = stall+1 if (abs(curL-prevL) < 0.0002 and abs(curR-prevR) < 0.0002) else 0
    prevL, prevR = curL, curR
    # A jaw that stops at a WIDE opening has hit the pot body, not a handle -- accepting that as
    # "grasped" is how an episode ends up dangling from one lid pinch while the other hand hangs
    # empty and the numeric check still reports success. Require a plausible handle-sized gap.
    if stall >= 8 and -args.grip_max < curL < -0.002 and -args.grip_max < curR < -0.002:
        print(f"[pot] both jaws stalled on handle-sized gaps: L={curL:.4f} R={curR:.4f} "
              f"(gaps {2*abs(curL)*100:.1f}/{2*abs(curR)*100:.1f} cm) after {k} steps", flush=True)
        break

_phase["v"] = "lift"
Ltop = Lhold+np.array([0, 0, args.lift], np.float32)
Rtop = Rhold+np.array([0, 0, args.lift], np.float32)
drive(Ltop, CLOSE, Rtop, CLOSE, 150)
_phase["v"] = "hold"
for _ in range(40):
    env.step(act2((Ltop+_CORRL["v"]).astype(np.float32), CLOSE, (Rtop+_CORRR["v"]).astype(np.float32), CLOSE))
    capture()

w = potw(); q = potq()
dz = float(w[2])-z0
tilt = tilt_from_rest(q)      # how far it rolled AWAY from how it was resting, not from world +z
_RESULT["v"] = "SUCCESS" if (dz > 0.06 and tilt < 25.0) else "FAIL"
_phase["v"] = "done"
print(f"[pot] EPISODE_RESULT: {_RESULT['v']} dz={dz:+.3f} tilt={tilt:.1f}deg "
      f"pot=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f}) start_z={z0:.3f}", flush=True)
for _ in range(16):
    capture()

os.makedirs(os.path.dirname(args.video), exist_ok=True)
# Drop the warm-up frames: before the renderer settles they come out with the wrong camera
# pose, unresolved textures and missing geometry.
if len(frames) > 6:
    frames = frames[2:]
if frames:
    imageio.mimsave(args.video, frames, fps=14)
print(f"[pot] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_POT_LIFT_OK", flush=True)