"""YAM BIMANUAL HANDOVER: the LEFT arm picks the basket up by one rim edge, passes it to the
RIGHT arm in the middle of the table, and the right arm carries it to the right-hand side and
sets it down.

This is the one task where both arms must cooperate on the SAME object: the receiving jaw has to
close on the opposite rim before the giving jaw opens, or the basket drops. Sequence:

    left approach -> left grasp rim -> left lift -> move to the handover pose ->
    right approach opposite rim -> right close -> left open -> left retreat ->
    right carry to the right side -> right lower -> right release

    python scripts/yam_handover.py --headless --video outputs/tasks/handover.mp4
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser()
parser.add_argument("--obj", default="rw_basket_dyn", help="object to hand over (registered via YAM_RW_OBJECTS)")
parser.add_argument("--start_xy", default="0.02,0.16", help="start x,y on the LEFT side (env-local)")
parser.add_argument("--handover_xy", default="0.06,0.00", help="where the pass happens (env-local)")
parser.add_argument("--goal_xy", default="0.02,-0.20", help="where the RIGHT arm sets it down")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/yam_handover.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 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 as _e:
    print("[ho] 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"[ho] {args.obj!r} not in scene -- register it with YAM_RW_OBJECTS")
OBJ = u.scene.rigid_objects[args.obj]
SXY = [float(v) for v in args.start_xy.split(",")]
HXY = [float(v) for v in args.handover_xy.split(",")]
GXY = [float(v) for v in args.goal_xy.split(",")]


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, _ = eef_root(L, Lbn, lroot, lrootq)
rp0, _ = eef_root(R, Rbn, rroot, rrootq)
GQ = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))


def act2(lp, lg, rp, rg):
    return torch.tensor(np.concatenate([lp, GQ, [lg], rp, GQ, [rg]]),
                        dtype=torch.float32, device=dev).view(1, -1)


def place_obj(xy, z):
    OBJ.write_root_pose_to_sim(torch.tensor(np.concatenate([origin+np.array([xy[0], xy[1], z]), [1, 0, 0, 0]]),
                                            dtype=torch.float32, device=dev).view(1, 7))
    OBJ.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))


place_obj(SXY, 0.60)
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(OBJ.root_physx_view.prim_paths[0])).ComputeAlignedRange()
ext = np.array(rng.GetMax())-np.array(rng.GetMin())
print(f"[ho] {args.obj} extents={np.round(ext,3)}", flush=True)
place_obj(SXY, TABLE_TOP+float(ext[2])/2.0+0.004)     # seat it on the table, no drop
for _ in range(80):
    env.step(act2(lp0, OPEN, rp0, OPEN))


def objw():
    return OBJ.data.root_pos_w[0].cpu().numpy()-origin


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


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


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


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"[ho] friction failed {tag}:", e, flush=True)


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

frames = []; _phase = {"v": "start"}; _RESULT = {"v": ""}; _G = {"l": "OPEN", "r": "OPEN"}


def capture():
    img = env.render()
    if img is None:
        return
    im = Image.fromarray(np.asarray(img)[..., :3].copy()); d = ImageDraw.Draw(im)
    w = objw()
    lines = ["=== BIMANUAL HANDOVER: basket left arm -> right arm ==="
             + (f"  EP {args.episode}" if args.episode >= 0 else "")]
    if _RESULT["v"]:
        lines.append(f"RESULT: {_RESULT['v']}")
    lines += [f"ACTION: {_phase['v']}",
              f"left={_G['l']}  right={_G['r']}   obj=({w[0]:+.2f},{w[1]:+.2f},{w[2]:.2f})"]
    d.rectangle([0, 0, 470, 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))


_CL = {"v": np.zeros(3, np.float32)}; _CR = {"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):
    """Move both arms along one eased profile; either target may be None to hold in place."""
    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 = ls if lt is None else np.asarray(lt, np.float32)
    rt = rs if rt is None else np.asarray(rt, np.float32)
    _G["l"] = "CLOSE" if lg < 0 else "OPEN"; _G["r"] = "CLOSE" if rg < 0 else "OPEN"
    cl, cr = _CL["v"], _CR["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"], _CMD["r"] = lc, rc
        env.step(act2((lc+cl).astype(np.float32), lg, (rc+cr).astype(np.float32), rg))
        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)
        _CL["v"], _CR["v"] = cl, cr
        if k % 3 == 0:
            capture()


def clamp(arm, lg, rg, n=150):
    """Close one jaw until it stalls on the rim; the other jaw holds whatever it is doing."""
    prev = fsep(arm); stall = 0
    for k in range(n):
        env.step(act2((_CMD["l"]+_CL["v"]).astype(np.float32), lg,
                      (_CMD["r"]+_CR["v"]).astype(np.float32), rg))
        if k % 5 == 0:
            capture()
        cur = fsep(arm)
        stall = stall+1 if abs(cur-prev) < 0.0002 else 0
        prev = cur
        if stall >= 8 and cur < -0.002:
            print(f"[ho] jaw stalled at fsep={cur:.4f} after {k} steps", flush=True)
            return True
    print(f"[ho] jaw did NOT stall (fsep={fsep(arm):.4f})", flush=True)
    return False


def to_L(w):
    return (Rq(lrootq).T@(np.asarray(w, np.float32)-lroot)).astype(np.float32)


def to_R(w):
    return (Rq(rrootq).T@(np.asarray(w, np.float32)-rroot)).astype(np.float32)


w0 = objw()
half_y = float(ext[1])/2.0
rim_z = TABLE_TOP+float(ext[2])*0.80          # grab the rim, near the top edge
z_start = float(objw()[2])
# left arm takes the +y rim, right arm will take the -y rim
L_grip_w = np.array([w0[0], w0[1]+half_y-0.004, rim_z], np.float32)
print(f"[ho] object at ({w0[0]:.3f},{w0[1]:.3f}) ext={np.round(ext,3)} rim_z={rim_z:.3f}", flush=True)

_phase["v"] = "1. LEFT ARM approach basket rim"
drive(to_L(L_grip_w+np.array([0, 0, 0.13], np.float32)), OPEN, None, OPEN, 130)
_phase["v"] = "2. LEFT descend to rim"
drive(to_L(L_grip_w), OPEN, None, OPEN, 110)
_phase["v"] = "3. LEFT close on rim"
clamp(L, CLOSE, OPEN)
_phase["v"] = "4. LEFT lift"
lift_w = np.array([w0[0], w0[1]+half_y-0.004, rim_z+0.12], np.float32)
drive(to_L(lift_w), CLOSE, None, OPEN, 130)
z_lift = float(objw()[2])
print(f"[ho] left lifted: obj z {z_start:.3f} -> {z_lift:.3f}", flush=True)

# carry to the handover pose in the middle, where BOTH arms can reach the object
_phase["v"] = "5. LEFT carry to handover pose"
ho = objw()
hand_w = np.array([HXY[0], HXY[1], TABLE_TOP+0.16], np.float32)
d_xy = hand_w[:2]-ho[:2]
L_hand = np.array([L_grip_w[0]+d_xy[0], L_grip_w[1]+d_xy[1], hand_w[2]], np.float32)
drive(to_L(L_hand), CLOSE, None, OPEN, 150)
oc = objw()
print(f"[ho] at handover: obj=({oc[0]:.3f},{oc[1]:.3f},{oc[2]:.3f})", flush=True)

_phase["v"] = "6. RIGHT arm approach opposite rim"
R_grip_w = np.array([oc[0], oc[1]-half_y+0.004, float(oc[2])+float(ext[2])*0.30], np.float32)
drive(None, CLOSE, to_R(R_grip_w+np.array([0, 0, 0.12], np.float32)), OPEN, 140)
_phase["v"] = "7. RIGHT descend onto rim"
drive(None, CLOSE, to_R(R_grip_w), OPEN, 110)
_phase["v"] = "8. RIGHT close (both hold)"
got = clamp(R, CLOSE, CLOSE)
_phase["v"] = "9. LEFT release"
drive(None, OPEN, None, CLOSE, 45)
_phase["v"] = "10. LEFT retreat"
drive(to_L(L_hand+np.array([0.0, 0.10, 0.10], np.float32)), OPEN, None, CLOSE, 110)
after = objw()
print(f"[ho] after handover: obj=({after[0]:.3f},{after[1]:.3f},{after[2]:.3f}) right_fsep={fsep(R):.4f}", flush=True)

_phase["v"] = "11. RIGHT carry to the right side"
goal_w = np.array([GXY[0], GXY[1], TABLE_TOP+float(ext[2])/2.0+0.02], np.float32)
d2 = goal_w[:2]-after[:2]
R_goal = np.array([R_grip_w[0]+d2[0], R_grip_w[1]+d2[1],
                   float(R_grip_w[2])+(goal_w[2]-after[2])], np.float32)
drive(None, OPEN, to_R(R_goal+np.array([0, 0, 0.10], np.float32)), CLOSE, 150)
_phase["v"] = "12. RIGHT lower onto the table"
drive(None, OPEN, to_R(R_goal), CLOSE, 110)
_phase["v"] = "13. RIGHT release"
drive(None, OPEN, None, OPEN, 45)
_phase["v"] = "14. RIGHT retreat"
drive(None, OPEN, to_R(R_goal+np.array([0, 0, 0.14], np.float32)), OPEN, 90)
for _ in range(60):
    env.step(act2((_CMD["l"]+_CL["v"]).astype(np.float32), OPEN, (_CMD["r"]+_CR["v"]).astype(np.float32), OPEN))

wf = objw()
handed = got and after[2] > TABLE_TOP+0.05        # right arm still held it up after left let go
on_right = wf[1] < GXY[1]+0.10 and wf[1] < 0.0    # ended on the right-hand side
settled = wf[2] < TABLE_TOP+float(ext[2])+0.03
_RESULT["v"] = "SUCCESS" if (handed and on_right and settled) else "FAIL"
_phase["v"] = "DONE"
print(f"[ho] EPISODE_RESULT: {_RESULT['v']} handed={handed} on_right={on_right} settled={settled} "
      f"final=({wf[0]:.3f},{wf[1]:.3f},{wf[2]:.3f}) goal=({GXY[0]:.2f},{GXY[1]:.2f})", 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"[ho] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_HANDOVER_OK", flush=True)