"""YAM BIMANUAL: take the lid off the pot, then put food inside.

The LEFT arm lifts the lid off the pot and parks it on the table; only then can the RIGHT arm
drop a piece of food into the open pot. The ordering is the point -- the food cannot go in while
the lid is on, so the two arms have to take turns on a shared workspace.

The pot is built from primitive walls and the lid is a separate rigid body, because RoboTwin's
060_kitchenpot is a single fused mesh whose lid cannot be removed.

    python scripts/yam_lid_food.py --headless --food apple --video outputs/tasks/lid_food.mp4
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser()
parser.add_argument("--food", default="apple", help="the item to put in the pot")
parser.add_argument("--food_xy", default="0.06,0.03", help="food start x,y (right side, env-local)")
parser.add_argument("--pot_xy", default="0.00,0.13", help="pot centre x,y (left side, env-local)")
parser.add_argument("--lid_park_xy", default="-0.10,0.20", help="where the left arm parks the lid")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/yam_lid_food.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("[lf] time_out disable failed:", _e)
try:
    _cfg.viewer.eye = (0.95, -0.95, 1.15); _cfg.viewer.lookat = (0.02, 0.05, 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
PXY = [float(v) for v in args.pot_xy.split(",")]
FXY = [float(v) for v in args.food_xy.split(",")]
KXY = [float(v) for v in args.lid_park_xy.split(",")]
for _n in ("lid", args.food):
    if _n not in u.scene.rigid_objects:
        raise SystemExit(f"[lf] {_n!r} not in scene")
LID = u.scene.rigid_objects["lid"]; FOOD = u.scene.rigid_objects[args.food]

# ---- the pot: four primitive walls leaving a rectangular mouth the lid sits on ----
import isaaclab.sim as sim_utils
MOUTH_X, MOUTH_Y, WALL_H, WALL_T = 0.135, 0.048, 0.085, 0.012


def _wall(name, size, off, color=(0.40, 0.42, 0.48)):
    c = sim_utils.CuboidCfg(size=tuple(size),
                            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color),
                            collision_props=sim_utils.CollisionPropertiesCfg())
    c.func(f"/World/envs/env_0/pot_{name}", c,
           translation=tuple((origin+np.array([PXY[0]+off[0], PXY[1]+off[1], TABLE_TOP+off[2]],
                                              np.float32)).astype(float).tolist()))


_wall("floor", (MOUTH_X+2*WALL_T, MOUTH_Y+2*WALL_T, 0.010), (0, 0, 0.005))
_wall("xp", (WALL_T, MOUTH_Y+2*WALL_T, WALL_H), (MOUTH_X/2+WALL_T/2, 0, WALL_H/2))
_wall("xn", (WALL_T, MOUTH_Y+2*WALL_T, WALL_H), (-MOUTH_X/2-WALL_T/2, 0, WALL_H/2))
_wall("yp", (MOUTH_X+2*WALL_T, WALL_T, WALL_H), (0, MOUTH_Y/2+WALL_T/2, WALL_H/2))
_wall("yn", (MOUTH_X+2*WALL_T, WALL_T, WALL_H), (0, -MOUTH_Y/2-WALL_T/2, WALL_H/2))
print(f"[lf] pot at ({PXY[0]},{PXY[1]}) mouth={MOUTH_X}x{MOUTH_Y} wall_h={WALL_H}", 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)
JAW_X = qR(np.stack([np.array([1., 0., 0.]), np.array([0., -1., 0.]), np.array([0., 0., -1.])], axis=1))
JAW_Y = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))


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)


LID_T = 0.016
LID.write_root_pose_to_sim(torch.tensor(
    np.concatenate([origin+np.array([PXY[0], PXY[1], TABLE_TOP+WALL_H+LID_T/2+0.002]), [1, 0, 0, 0]]),
    dtype=torch.float32, device=dev).view(1, 7))
LID.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))
FOOD.write_root_pose_to_sim(torch.tensor(
    np.concatenate([origin+np.array([FXY[0], FXY[1], 0.52]), [1, 0, 0, 0]]),
    dtype=torch.float32, device=dev).view(1, 7))
FOOD.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))
for _ in range(90):
    env.step(act2(lp0, JAW_X, OPEN, rp0, JAW_Y, OPEN))


def _boost(view, tag, s=1.7, d=1.5):
    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"[lf] friction failed {tag}:", e, flush=True)


_boost(R.root_physx_view, "right"); _boost(L.root_physx_view, "left")
_boost(LID.root_physx_view, "lid"); _boost(FOOD.root_physx_view, args.food)

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])
frng = bbc.ComputeWorldBound(stage.GetPrimAtPath(FOOD.root_physx_view.prim_paths[0])).ComputeAlignedRange()
fext = np.array(frng.GetMax())-np.array(frng.GetMin())
print(f"[lf] food {args.food} size={np.round(fext,3)}", flush=True)


def lidw():
    return LID.data.root_pos_w[0].cpu().numpy()-origin


def foodw():
    return FOOD.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


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)
    lw, fw = lidw(), foodw()
    lines = ["=== LID OFF THE POT, THEN FOOD IN ===" + (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']}",
              f"lid=({lw[0]:+.2f},{lw[1]:+.2f},{lw[2]:.2f})  food=({fw[0]:+.2f},{fw[1]:+.2f},{fw[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}
_Q = {"l": JAW_X, "r": JAW_Y}


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


def drive(lt, rt, lg, rg, 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["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), _Q["l"], lg, (rc+cr).astype(np.float32), _Q["r"], 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):
    prev = fsep(arm); stall = 0
    for k in range(n):
        env.step(act2((_CMD["l"]+_CL["v"]).astype(np.float32), _Q["l"], lg,
                      (_CMD["r"]+_CR["v"]).astype(np.float32), _Q["r"], 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"[lf] jaw stalled at fsep={cur:.4f} after {k}", flush=True)
            return True
    print(f"[lf] 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)


# ---------------- phase 1: LEFT arm takes the lid off ----------------
lid0 = lidw()
lid_grip = np.array([lid0[0], lid0[1], TABLE_TOP+WALL_H+LID_T*0.4], np.float32)
print(f"[lf] lid at ({lid0[0]:.3f},{lid0[1]:.3f},{lid0[2]:.3f}) grip_z={lid_grip[2]:.3f}", flush=True)
_phase["v"] = "1. LEFT approach the lid"
drive(to_L(lid_grip+np.array([0, 0, 0.13], np.float32)), None, OPEN, OPEN, 130)
_phase["v"] = "2. LEFT descend onto the lid"
drive(to_L(lid_grip), None, OPEN, OPEN, 110)
got_lid = clamp(L, CLOSE, OPEN)
_phase["v"] = "3. LEFT lift the lid clear"
drive(to_L(lid_grip+np.array([0, 0, 0.15], np.float32)), None, CLOSE, OPEN, 130)
z_lid = float(lidw()[2])
_phase["v"] = "4. LEFT park the lid aside"
park = np.array([KXY[0], KXY[1], TABLE_TOP+0.15], np.float32)
drive(to_L(park), None, CLOSE, OPEN, 150)
drive(to_L(np.array([KXY[0], KXY[1], TABLE_TOP+LID_T/2+0.02], np.float32)), None, CLOSE, OPEN, 110)
_phase["v"] = "5. LEFT release the lid"
drive(None, None, OPEN, OPEN, 45)
drive(to_L(park+np.array([0, 0, 0.06], np.float32)), None, OPEN, OPEN, 90)
lid_off = float(np.hypot(lidw()[0]-PXY[0], lidw()[1]-PXY[1])) > 0.08
print(f"[lf] lid removed={lid_off} lifted_to={z_lid:.3f} now=({lidw()[0]:.3f},{lidw()[1]:.3f})", flush=True)

# ---------------- phase 2: RIGHT arm puts the food in ----------------
fw = foodw()
food_grip = np.array([fw[0], fw[1], TABLE_TOP+float(fext[2])/2.0], np.float32)
_phase["v"] = "6. RIGHT approach the food"
drive(None, to_R(food_grip+np.array([0, 0, 0.13], np.float32)), OPEN, OPEN, 130)
_phase["v"] = "7. RIGHT descend onto the food"
drive(None, to_R(food_grip), OPEN, OPEN, 120)
got_food = clamp(R, OPEN, CLOSE)
f_z0 = float(foodw()[2])
_phase["v"] = "8. RIGHT lift the food"
drive(None, to_R(food_grip+np.array([0, 0, 0.17], np.float32)), OPEN, CLOSE, 130)
_phase["v"] = "9. RIGHT carry over the open pot"
over = np.array([PXY[0], PXY[1], TABLE_TOP+WALL_H+float(fext[2])/2.0+0.05], np.float32)
drive(None, to_R(over+np.array([0, 0, 0.06], np.float32)), OPEN, CLOSE, 160)
_phase["v"] = "10. RIGHT lower into the pot"
drive(None, to_R(over), OPEN, CLOSE, 110)
_phase["v"] = "11. RIGHT release the food"
drive(None, None, OPEN, OPEN, 45)
_phase["v"] = "12. RIGHT retreat"
drive(None, to_R(over+np.array([0, 0, 0.14], np.float32)), OPEN, OPEN, 100)
for _ in range(70):
    env.step(act2((_CMD["l"]+_CL["v"]).astype(np.float32), _Q["l"], OPEN,
                  (_CMD["r"]+_CR["v"]).astype(np.float32), _Q["r"], OPEN))

lf = lidw(); ff = foodw()
food_in = (abs(ff[0]-PXY[0]) < MOUTH_X/2+0.02 and abs(ff[1]-PXY[1]) < MOUTH_Y/2+0.03
           and ff[2] < TABLE_TOP+WALL_H)
_RESULT["v"] = "SUCCESS" if (lid_off and food_in) else "FAIL"
_phase["v"] = "DONE"
print(f"[lf] EPISODE_RESULT: {_RESULT['v']} lid_off={lid_off} food_in_pot={food_in} "
      f"lid=({lf[0]:.3f},{lf[1]:.3f},{lf[2]:.3f}) food=({ff[0]:.3f},{ff[1]:.3f},{ff[2]:.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"[lf] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_LID_FOOD_OK", flush=True)