"""v4: render BOTH arms (left in green, right in red) using T_left_from_right."""
import mujoco, numpy as np, pickle, cv2, sys

MJCF = "/home/robot-lab/raiden/third_party/i2rt/i2rt/robot_models/arm/yam/yam.xml"
DATA = "/home/robot-lab/data/processed/pickup_apple/0000"
FRAME = int(sys.argv[1]) if len(sys.argv) > 1 else 500
CAM = sys.argv[2] if len(sys.argv) > 2 else "scene_1"
OUT = f"/home/robot-lab/cameron/yam_overlay/out_v4_{CAM}_f{FRAME:04d}.png"

fstr = f"{FRAME:010d}"

model = mujoco.MjModel.from_xml_path(MJCF)
data = mujoco.MjData(model)
with open(f"{DATA}/lowdim/{fstr}.pkl", "rb") as f:
    fd = pickle.load(f)

# LEFT ARM positions in left_arm_base (= world) frame
data.qpos[:6] = fd["joints"][:6]
mujoco.mj_kinematics(model, data)
left_world = data.xpos.copy()

# RIGHT ARM positions in right_arm_base frame
data.qpos[:6] = fd["joints"][7:13]
mujoco.mj_kinematics(model, data)
right_in_right_frame = data.xpos.copy()

# Transform right arm positions to left_arm_base (= world) frame
T_left_from_right = np.array(fd["T_left_from_right"])
print(f"T_left_from_right =\n{T_left_from_right}")
print(f"T_left_from_right translation: {T_left_from_right[:3, 3]}")

pts_h = np.hstack([right_in_right_frame, np.ones((model.nbody, 1))])
right_world = (T_left_from_right @ pts_h.T).T[:, :3]

print(f"\nLEFT arm world positions (link_6 = gripper):")
for i in range(model.nbody):
    print(f"  L_{model.body(i).name}: {left_world[i]}")
print(f"\nRIGHT arm world positions (in left_arm_base frame):")
for i in range(model.nbody):
    print(f"  R_{model.body(i).name}: {right_world[i]}")

K = np.array(fd["intrinsics"][CAM])
T_cam2world = np.array(fd["extrinsics"][CAM])
T_world2cam = np.linalg.inv(T_cam2world)
img = cv2.imread(f"{DATA}/rgb/{CAM}/{fstr}.jpg")
H, W = img.shape[:2]

def draw(pts_world, color, prefix):
    pts_h = np.hstack([pts_world, np.ones((pts_world.shape[0], 1))])
    pts_cam = (T_world2cam @ pts_h.T).T[:, :3]
    mask = pts_cam[:, 2] > 0.01
    denom = np.where(pts_cam[:, 2:3] != 0, pts_cam[:, 2:3], 1)
    uv = (K @ (pts_cam / denom).T).T[:, :2]
    print(f"\n[{prefix}] uv:")
    for i in range(pts_world.shape[0]):
        u, v = uv[i]
        ok = mask[i] and 0 <= u < W and 0 <= v < H
        print(f"  {prefix}_{model.body(i).name}: ({u:.0f}, {v:.0f}) {'OK' if ok else ('OOB' if mask[i] else 'BEHIND')}")
    for i in range(pts_world.shape[0]):
        if not mask[i]: continue
        u, v = uv[i]
        if not (0 <= u < W and 0 <= v < H): continue
        cv2.circle(img, (int(u), int(v)), 8, color, -1)
        cv2.putText(img, f"{prefix}_{model.body(i).name}", (int(u)+8, int(v)-4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1, cv2.LINE_AA)
    for i in range(1, pts_world.shape[0]):
        p = model.body_parentid[i]
        if not (mask[i] and mask[p]): continue
        p0, p1 = (int(uv[p,0]), int(uv[p,1])), (int(uv[i,0]), int(uv[i,1]))
        if all(0 <= c[0] < W and 0 <= c[1] < H for c in (p0, p1)):
            cv2.line(img, p0, p1, color, 2)

draw(left_world, (0, 255, 0), "L")     # green = left arm (joints[:6])
draw(right_world, (0, 0, 255), "R")    # red = right arm (joints[7:13])

cv2.putText(img, f"L=green(joints[:6])  R=red(joints[7:13])", (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)

cv2.imwrite(OUT, img)
print(f"\nSaved: {OUT}")