"""YAM BIMANUAL plank carry: a long plank lies on the table at an ANGLE, and a target region is
marked on the table at its own angle. Both arms grip the plank's two ends, lift it, carry it to
the target and set it down aligned with the marked region.

One gripper cannot span a 36 cm plank, so both arms must hold it and move on the same profile;
and because the plank is diagonal, the grip points are computed along its real axis rather than
along x/y, and each wrist is yawed to match.

    python scripts/yam_plank_place.py --headless --plank_yaw 25 --place_yaw -20 \
        --video outputs/tasks/plank.mp4
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser()
parser.add_argument("--obj", default="board")
parser.add_argument("--plank_xy", default="0.06,0.02", help="plank centre x,y (env-local)")
parser.add_argument("--plank_yaw", type=float, default=25.0, help="plank yaw on the table (deg)")
parser.add_argument("--place_xy", default="0.06,-0.20", help="target region centre x,y")
parser.add_argument("--place_yaw", type=float, default=-20.0, help="target region yaw (deg)")
parser.add_argument("--lift", type=float, default=0.13)
parser.add_argument("--grip_inset", type=float, default=0.04, help="grip this far in from each end")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/yam_plank_place.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("[pl] 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


def yaw_quat(deg):
    a = np.radians(deg)
    return np.array([np.cos(a/2), 0.0, 0.0, np.sin(a/2)])


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
PXY = [float(v) for v in args.plank_xy.split(",")]
QXY = [float(v) for v in args.place_xy.split(",")]
OBJ = u.scene.rigid_objects[args.obj]

# ---- target region marker, drawn at the requested angle ----
import isaaclab.sim as sim_utils
rng_ext = None
_m = sim_utils.CuboidCfg(size=(0.40, 0.09, 0.0015),
                         visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.20, 0.65, 0.30)))
_m.func("/World/envs/env_0/plank_target", _m,
        translation=tuple((origin+np.array([QXY[0], QXY[1], TABLE_TOP+0.001])).astype(float).tolist()),
        orientation=tuple(float(v) for v in yaw_quat(args.place_yaw)))
print(f"[pl] target region at ({QXY[0]},{QXY[1]}) yaw={args.place_yaw} deg", flush=True)


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)


def grasp_quat(yaw_deg):
    """Top-down grasp whose jaw closes ACROSS the plank, i.e. perpendicular to its axis."""
    a = np.radians(yaw_deg)
    ca, sa = np.cos(a), np.sin(a)
    Rz = np.array([[ca, -sa, 0], [sa, ca, 0], [0, 0, 1]])
    base = np.stack([np.array([1., 0., 0.]), np.array([0., -1., 0.]), np.array([0., 0., -1.])], axis=1)
    return qR(Rz@base)


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


GQ0 = grasp_quat(0.0)


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


seat_plank(args.plank_yaw, 0.60)
for _ in range(60):
    env.step(act2(lp0, GQ0, OPEN, rp0, GQ0, 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())
LEN = float(max(ext[0], ext[1])); THK = float(ext[2])
seat_plank(args.plank_yaw, TABLE_TOP+THK/2.0+0.003)     # seat it flat, no drop
for _ in range(90):
    env.step(act2(lp0, GQ0, OPEN, rp0, GQ0, OPEN))
print(f"[pl] plank len={LEN:.3f} thickness={THK:.3f} yaw={args.plank_yaw} deg", 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"[pl] 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)


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


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


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


def obj_yaw():
    q = OBJ.data.root_quat_w[0].cpu().numpy()
    w, x, y, z = [float(v) for v in q]
    return float(np.degrees(np.arctan2(2*(w*z+x*y), 1-2*(y*y+z*z))))


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


frames = []; _phase = {"v": "start"}; _RESULT = {"v": ""}; _G = {"v": "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 PLANK CARRY (angled) ===" + (f"  EP {args.episode}" if args.episode >= 0 else "")]
    if _RESULT["v"]:
        lines.append(f"RESULT: {_RESULT['v']}")
    lines += [f"ACTION: {_phase['v']}",
              f"grippers={_G['v']}  plank=({w[0]:+.2f},{w[1]:+.2f},{w[2]:.2f}) yaw={obj_yaw():+.0f}deg",
              f"target=({QXY[0]:+.2f},{QXY[1]:+.2f}) yaw={args.place_yaw:+.0f}deg"]
    d.rectangle([0, 0, 450, 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, rt, lq, rq, g, n):
    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["v"] = "CLOSE" if g < 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), lq, g, (rc+cr).astype(np.float32), rq, g))
        L.write_joint_state_to_sim(lhome_q.view(1, -1), torch.zeros((1, lhome_q.shape[0]), device=dev)) if False else None
        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_both(lq, rq, n=160):
    prevL, prevR = fsep(L), fsep(R); stall = 0
    for k in range(n):
        env.step(act2((_CMD["l"]+_CL["v"]).astype(np.float32), lq, CLOSE,
                      (_CMD["r"]+_CR["v"]).astype(np.float32), rq, 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
        if stall >= 8 and curL < -0.002 and curR < -0.002:
            print(f"[pl] both jaws stalled L={curL:.4f} R={curR:.4f} after {k}", flush=True)
            return 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)


def ends(centre_xy, yaw_deg, z):
    """The two grip points along the plank's own axis."""
    a = np.radians(yaw_deg)
    d = np.array([np.cos(a), np.sin(a)])
    r = LEN/2.0-args.grip_inset
    p_plus = np.array([centre_xy[0]+d[0]*r, centre_xy[1]+d[1]*r, z], np.float32)
    p_minus = np.array([centre_xy[0]-d[0]*r, centre_xy[1]-d[1]*r, z], np.float32)
    # the +y end belongs to the left arm
    return (p_plus, p_minus) if p_plus[1] > p_minus[1] else (p_minus, p_plus)


w0 = objw()
grip_z = TABLE_TOP+THK*0.55
LG, RG = ends((w0[0], w0[1]), args.plank_yaw, grip_z)
GQ = grasp_quat(args.plank_yaw)
print(f"[pl] grips: L={np.round(LG,3)} R={np.round(RG,3)} grip_z={grip_z:.3f}", flush=True)

_phase["v"] = "1. BOTH ARMS APPROACH plank ends"
drive(to_L(LG+np.array([0, 0, 0.13], np.float32)), to_R(RG+np.array([0, 0, 0.13], np.float32)), GQ, GQ, OPEN, 130)
_phase["v"] = "2. DESCEND onto the ends"
drive(to_L(LG), to_R(RG), GQ, GQ, OPEN, 110)
print(f"[pl] descended L_err={np.linalg.norm(eefL()-to_L(LG)):.3f} R_err={np.linalg.norm(eefR()-to_R(RG)):.3f}", flush=True)
_phase["v"] = "3. BOTH JAWS CLOSE"
got = clamp_both(GQ, GQ)
z0 = float(objw()[2])
_phase["v"] = "4. SYNCHRONISED LIFT"
drive(to_L(LG+np.array([0, 0, args.lift], np.float32)), to_R(RG+np.array([0, 0, args.lift], np.float32)),
      GQ, GQ, CLOSE, 150)
z_lift = float(objw()[2])
print(f"[pl] lifted: z {z0:.3f} -> {z_lift:.3f}", flush=True)

# carry to the target region and rotate to its angle: recompute the two grip points at the
# target pose and drive each wrist there, yawing the jaws to match
_phase["v"] = "5. CARRY + ROTATE to the target angle"
TL, TR = ends((QXY[0], QXY[1]), args.place_yaw, grip_z+args.lift)
TQ = grasp_quat(args.place_yaw)
drive(to_L(TL), to_R(TR), TQ, TQ, CLOSE, 190)
_phase["v"] = "6. LOWER onto the target region"
TL2, TR2 = ends((QXY[0], QXY[1]), args.place_yaw, grip_z+0.006)
drive(to_L(TL2), to_R(TR2), TQ, TQ, CLOSE, 130)
_phase["v"] = "7. RELEASE"
drive(None, None, TQ, TQ, OPEN, 50)
_phase["v"] = "8. RETREAT"
drive(to_L(TL2+np.array([0, 0, 0.15], np.float32)), to_R(TR2+np.array([0, 0, 0.15], np.float32)),
      TQ, TQ, OPEN, 110)
for _ in range(60):
    env.step(act2((_CMD["l"]+_CL["v"]).astype(np.float32), TQ, OPEN,
                  (_CMD["r"]+_CR["v"]).astype(np.float32), TQ, OPEN))

wf = objw(); yf = obj_yaw()
d_xy = float(np.hypot(wf[0]-QXY[0], wf[1]-QXY[1]))
d_yaw = abs(((yf-args.place_yaw)+90) % 180-90)      # plank is symmetric: 180 deg is the same pose
lifted = (z_lift-z0) > 0.05
_RESULT["v"] = "SUCCESS" if (lifted and d_xy < 0.06 and d_yaw < 18.0) else "FAIL"
_phase["v"] = "DONE"
print(f"[pl] EPISODE_RESULT: {_RESULT['v']} lifted={lifted} dz={z_lift-z0:+.3f} "
      f"pos_err={d_xy:.3f} yaw_err={d_yaw:.1f}deg final=({wf[0]:.3f},{wf[1]:.3f},{wf[2]:.3f}) yaw={yf:.1f}", 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"[pl] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_PLANK_OK", flush=True)