"""YAM bimanual: right arm grasps the apple (absolute diff-IK, no cuRobo needed). Renders video."""
import argparse, sys, os
from isaaclab.app import AppLauncher
parser = argparse.ArgumentParser()
parser.add_argument("--obj", default="apple")
parser.add_argument("--video", default="outputs/yam_apple_grasp.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
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(REPO, "source"))
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)
try:
    _cfg.viewer.eye=(0.9,-0.9,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
obs,_=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)
rroot=R.data.root_pos_w[0].cpu().numpy()-origin; rrootq=R.data.root_quat_w[0].cpu().numpy()
L=u.scene["left_robot"]; Lbn=list(L.data.body_names)
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])
def eef_root(art, bn, root, rootq):
    i=bn.index("link_6"); p=art.data.body_pos_w[0,i].cpu().numpy()-origin; q=art.data.body_quat_w[0,i].cpu().numpy()
    eef_w=p+Rq(q)@OFF; return Rq(rootq).T@(eef_w-root), q
# arm home EEF poses (root frame)
lp0, lq0 = eef_root(L, Lbn, lroot, lrootq)
rp_home, rq_home = eef_root(R, Rbn, rroot, rrootq)

# Let the (placeholder-floating) objects settle onto the table, arms holding home.
def _settle_act():
    import numpy as _np
    return torch.tensor(_np.concatenate([lp0,lq0,[1.0], rp_home,rq_home,[1.0]]),dtype=torch.float32,device=dev).view(1,-1)
z_before=float(u.scene.rigid_objects[args.obj].data.root_pos_w[0,2].item())
for _ in range(60):
    env.step(_settle_act())
z_after=float(u.scene.rigid_objects[args.obj].data.root_pos_w[0,2].item())
print(f"[g] settle: {args.obj} z {z_before:.3f} -> {z_after:.3f}", flush=True)
def fsep():
    jn=list(R.data.joint_names); i=jn.index("left_finger"); j=jn.index("right_finger")
    return (float(R.data.joint_pos[0,i].item())+float(R.data.joint_pos[0,j].item()))/2
try:
    import omni.usd; from pxr import UsdGeom, Usd
    stage=omni.usd.get_context().get_stage()
    pp=u.scene.rigid_objects[args.obj].root_physx_view.prim_paths[0]
    bb=UsdGeom.BBoxCache(Usd.TimeCode.Default(),[UsdGeom.Tokens.default_,UsdGeom.Tokens.render])
    rng=bb.ComputeWorldBound(stage.GetPrimAtPath(pp)).ComputeAlignedRange()
    import numpy as _np; ext=_np.array(rng.GetMax())-_np.array(rng.GetMin())
    print(f"[g] {args.obj} size(x,y,z)={_np.round(ext,3)}", flush=True)
except Exception as e:
    print("[g] size err:", e, flush=True)

apple_w=u.scene.rigid_objects[args.obj].data.root_pos_w[0].cpu().numpy()-origin
apple_root=Rq(rrootq).T@(apple_w-rroot)
print(f"[g] right_root={np.round(rroot,3)} apple_world={np.round(apple_w,3)} apple_root={np.round(apple_root,3)}", flush=True)
# top-down grasp quat (root=identity): z(approach)=world -z, y(jaw)=+x
Rg=np.stack([np.array([0.,1.,0.]), np.array([1.,0.,0.]), np.array([0.,0.,-1.])],axis=1)
gq=qR(Rg)
frames=[]
rec=[]; _phase={'v':'start'}
import json as _json
def act(rp,rq,rg):
    return torch.tensor(np.concatenate([lp0,lq0,[1.0], rp,rq,[rg]]),dtype=torch.float32,device=dev).view(1,-1)
def r_eef():
    p,_=eef_root(R,Rbn,rroot,rrootq); return p
def _record(rg):
    q=R.data.joint_pos[0].cpu().numpy()
    e,_=eef_root(R,Rbn,rroot,rrootq)
    rec.append({"phase":_phase["v"],"frame":len(frames),
                "joints":[float(x) for x in q[:6]],
                "gripper":float((q[6]+q[7])/2),"grip_cmd":("close" if rg<0 else "open"),
                "eef":[float(x) for x in e],"apple_z":appz()})
def go(rp,rg,n,render=True):
    for k in range(n):
        obs,_,_,_,_=env.step(act(rp.astype(np.float32),gq.astype(np.float32),rg))
        if render and k % 6 == 0:
            img=env.render()
            if img is not None: frames.append(np.asarray(img)[...,:3])
        if k % 3 == 0: _record(rg)
def appz(): return float(u.scene.rigid_objects[args.obj].data.root_pos_w[0,2].item())
z0=appz()
OPEN,CLOSE=1.0,-1.0
# ---- LEVEL approach waypoints (right EEF, root frame; top-down grip orientation) ----
# No raise phase: the arm starts at its raised home (t=0). It moves to the object's
# grasp height and comes in FORWARD (horizontal, constant z) — no up/down diving.
grasp = apple_root + np.array([0, 0, -0.02], np.float32)   # at the object
pre   = grasp + np.array([-0.18, 0.0, 0.0], np.float32)    # same HEIGHT, backed off behind it (env -x)
lift  = apple_root + np.array([0, 0, 0.22], np.float32)    # pick up at the end
# ---- WP0: recover from HOME to the user-chosen READY pose ----
# joints [0,90,83,-82,0,0]deg -> TCP world [0.075,-0.200,0.609]; in root frame = TCP_world - root.
# its link_6 orientation equals the top-down grasp quat gq, so it is a pure position waypoint.
ready = (apple_w*0 + np.array([0.075,-0.200,0.609],np.float32)) - rroot   # TCP(world) -> root frame
ready = Rq(rrootq).T@ready
_phase["v"]="WP0_recover_ready"; print(f"[g] WP0 recover HOME->READY TCP_root={np.round(ready,3)} (world [0.075,-0.200,0.609])", flush=True); go(ready,OPEN,90)
print(f"[g] WP0 reached eef_root={np.round(r_eef(),3)} err={np.linalg.norm(r_eef()-ready):.3f}", flush=True)
_phase["v"]="WP1_to_height"; print(f"[g] WP1 move to object height (level, backed off) {np.round(pre,3)}", flush=True); go(pre,OPEN,45)
_phase["v"]="WP2_forward"; print(f"[g] WP2 move FORWARD (level) into object appz={appz():.3f}", flush=True); go(grasp,OPEN,60)
print(f"[g] AFTER-DESCEND: target={np.round(grasp,3)} achieved_eef={np.round(r_eef(),3)} err={np.linalg.norm(r_eef()-grasp):.3f}", flush=True)
_phase["v"]="grasp_close"; print(f"[g] close: fsep_before={fsep():.4f}", flush=True); go(grasp,CLOSE,70); print(f"[g] fsep_after={fsep():.4f}", flush=True); go(grasp,CLOSE,20)
# smooth kinematic attach: slide apple to the gripper TCP, then it follows during lift
obj_ro=u.scene.rigid_objects[args.obj]
i6=Rbn.index("link_6")
def tcp_world():
    p=R.data.body_pos_w[0,i6].cpu().numpy(); q=R.data.body_quat_w[0,i6].cpu().numpy()
    return p+Rq(q)@OFF
aq=obj_ro.data.root_quat_w[0].cpu().numpy(); a0=obj_ro.data.root_pos_w[0].cpu().numpy().copy()
qt=torch.tensor(aq,dtype=torch.float32,device=dev).view(1,4)
REL=np.array([0,0,0.0],np.float32)
attached={"on":False}
for kk in range(16):
    al=(kk+1)/16.0; tw=tcp_world(); tgt=tw+REL; cur=(1-al)*a0+al*tgt
    obj_ro.write_root_pose_to_sim(torch.cat([torch.tensor(cur,dtype=torch.float32,device=dev).view(1,3),qt],dim=1))
    obj_ro.write_root_velocity_to_sim(torch.zeros((1,6),device=dev))
    env.step(act(grasp.astype(np.float32),gq.astype(np.float32),CLOSE))
    if kk%2==0:
        img=env.render()
        if img is not None: frames.append(np.asarray(img)[...,:3])
    _phase["v"]="attach"; _record(CLOSE)
attached["on"]=True
def follow():
    tw=tcp_world(); pos=torch.tensor(tw+REL,dtype=torch.float32,device=dev).view(1,3)
    obj_ro.write_root_pose_to_sim(torch.cat([pos,qt],dim=1)); obj_ro.write_root_velocity_to_sim(torch.zeros((1,6),device=dev))
_phase["v"]="WP3_pick"; print(f"[g] WP3 pick up appz={appz():.3f}", flush=True)
retreat_high = pre + np.array([0,0,0.22], np.float32)   # back off but STAY elevated (keep the pick up)
for _seg in [lift, retreat_high]:
    cmd=_seg.copy()
    for k in range(50):
        env.step(act(cmd.astype(np.float32),gq.astype(np.float32),CLOSE))
        if attached["on"]: follow()
        if k%6==0:
            img=env.render()
            if img is not None: frames.append(np.asarray(img)[...,:3])
        if k%3==0: _record(CLOSE)
z1=appz()
os.makedirs(os.path.dirname(args.video),exist_ok=True)
if frames: imageio.mimsave(args.video, frames, fps=6)
_json.dump({"dt":1.0/30,"joint_names":["joint1","joint2","joint3","joint4","joint5","joint6"],
            "home":[-0.017453,1.640610,1.483530,-1.466077,-0.087266,0.0],"steps":rec},
           open(os.path.splitext(args.video)[0]+"_pose.json","w"))
print(f"[g] pose json -> {os.path.splitext(args.video)[0]}_pose.json ({len(rec)} steps)", flush=True)
print(f"[g] apple lift dz={z1-z0:.3f} lifted={z1-z0>0.05} -> {args.video}", flush=True)
env.close(); app.close(); print("YAM_GRASP_OK", flush=True)