"""v9: render full-res overlay AND crop tight ROIs (300x300) around each marker
so we can inspect pixel-level alignment on the gripper.
"""
import numpy as np, pickle, cv2, sys, os

TASK = sys.argv[1]
EP = sys.argv[2]
CAM = sys.argv[3]
FRAME = int(sys.argv[4])

DATA = f"/home/robot-lab/data/processed/{TASK}/{EP}"
OUT_DIR = "/home/robot-lab/cameron/yam_overlay"

fstr = f"{FRAME:010d}"
pkl_path = f"{DATA}/lowdim/{fstr}.pkl"
img_path = None
for ext in (".jpg", ".png", ".jpeg"):
    cand = f"{DATA}/rgb/{CAM}/{fstr}{ext}"
    if os.path.exists(cand):
        img_path = cand; break
assert img_path, f"no image for {fstr}"

with open(pkl_path, "rb") as f:
    fd = pickle.load(f)

action = fd["action"]
l_pos = action[0:3]
r_pos_rb = action[13:16]
T_lfr = np.array(fd["T_left_from_right"])
r_pos_world = (T_lfr @ np.append(r_pos_rb, 1.0))[:3]
r_base_world = T_lfr[:3, 3]

K = np.array(fd["intrinsics"][CAM])
T_cam2world = np.array(fd["extrinsics"][CAM])
T_world2cam = np.linalg.inv(T_cam2world)
img_orig = cv2.imread(img_path)
H, W = img_orig.shape[:2]

def project(p_world):
    p_cam = (T_world2cam @ np.append(p_world, 1.0))[:3]
    if p_cam[2] < 0.01: return None
    uv = K @ (p_cam / p_cam[2])
    return float(uv[0]), float(uv[1])

points = {
    "LEFT_EE":   (l_pos,        (0, 255, 0)),
    "RIGHT_EE":  (r_pos_world,  (0, 0, 255)),
    "RIGHT_BASE":(r_base_world, (255, 0, 255)),
}

# Crop a 300x300 ROI centered on each marker, then draw a small + on it
HALF = 150
for name, (p_w, color) in points.items():
    uv = project(p_w)
    if uv is None: continue
    u, v = uv
    u_i, v_i = int(round(u)), int(round(v))
    x0, y0 = max(0, u_i - HALF), max(0, v_i - HALF)
    x1, y1 = min(W, u_i + HALF), min(H, v_i + HALF)
    roi = img_orig[y0:y1, x0:x1].copy()
    if roi.size == 0:
        print(f"  {name}: OOB, no ROI"); continue
    # Draw the marker in ROI coordinates
    cu, cv_ = u_i - x0, v_i - y0
    cv2.drawMarker(roi, (cu, cv_), color, cv2.MARKER_CROSS, 40, 2)
    cv2.circle(roi, (cu, cv_), 8, color, 2)
    cv2.putText(roi, name, (5, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
    cv2.putText(roi, f"uv=({u:.0f},{v:.0f})", (5, roi.shape[0]-10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1, cv2.LINE_AA)
    out = f"{OUT_DIR}/crop_{TASK}_{EP}_{CAM}_f{FRAME:04d}_{name}.png"
    cv2.imwrite(out, roi)
    print(f"  {name}: saved {out} (uv=({u:.1f},{v:.1f}))")