"""v7: same as v6, but invert T_left_from_right for the RIGHT arm.
Hypothesis: T_left_from_right name lies; the stored matrix maps left->right
(so its inverse is what we need to bring right_arm_base coords into world/left_arm_base).
"""
import numpy as np, pickle, cv2, sys

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_v7_{CAM}_f{FRAME:04d}.png"

fstr = f"{FRAME:010d}"
with open(f"{DATA}/lowdim/{fstr}.pkl", "rb") as f:
    fd = pickle.load(f)

action = fd["action"]
l_pos = action[0:3]
r_pos_rb = action[13:16]
T_left_from_right = np.array(fd["T_left_from_right"])
T_inv = np.linalg.inv(T_left_from_right)
r_pos_v6 = (T_left_from_right @ np.append(r_pos_rb, 1.0))[:3]
r_pos_v7 = (T_inv @ np.append(r_pos_rb, 1.0))[:3]

print(f"l_pos (world): {l_pos}")
print(f"r_pos in right_base: {r_pos_rb}")
print(f"r_pos via T_left_from_right (v6):     {r_pos_v6}")
print(f"r_pos via inv(T_left_from_right) (v7): {r_pos_v7}")
print(f"T_left_from_right translation: {T_left_from_right[:3,3]}")
print(f"inv translation:               {T_inv[:3,3]}")
print(f"l_grip: {action[12]}, r_grip: {action[25]}")

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 proj_one(p_world, label, color):
    p_h = np.append(p_world, 1.0)
    p_cam = (T_world2cam @ p_h)[:3]
    print(f"  {label}: world={p_world}, cam={p_cam}", end=" ")
    if p_cam[2] < 0.01:
        print("BEHIND CAM"); return
    uv = K @ (p_cam / p_cam[2])
    u, v = uv[0], uv[1]
    print(f"uv=({u:.1f}, {v:.1f})", end=" ")
    if 0 <= u < W and 0 <= v < H:
        cv2.drawMarker(img, (int(u), int(v)), color, cv2.MARKER_CROSS, 60, 4)
        cv2.circle(img, (int(u), int(v)), 12, color, 3)
        cv2.putText(img, label, (int(u)+20, int(v)-12),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2, cv2.LINE_AA)
        print("DRAWN")
    else:
        print("OOB")

proj_one(l_pos,    "LEFT_EE",         (0, 255, 0))    # green
proj_one(r_pos_v7, "RIGHT_EE_v7_inv", (0, 0, 255))    # red (new, inverted)
proj_one(r_pos_v6, "RIGHT_EE_v6",     (0, 165, 255))  # orange (old, for comparison)

cv2.putText(img, f"frame {FRAME}  v7: inv(T_lfr) for right", (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)

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