"""v8: validate v6 convention (raiden-canonical: T_left_from_right @ r_pos_rb, no inversion)
on a right-arm-active task across multiple frames. Also project right_arm_base origin
(= T_left_from_right translation) as a sanity marker.

Usage: python v8_multiframe.py <task> <episode> <cam> <frame1> [<frame2> ...]
"""
import numpy as np, pickle, cv2, sys, os

TASK = sys.argv[1]
EP = sys.argv[2]
CAM = sys.argv[3]
FRAMES = [int(x) for x in sys.argv[4:]]

DATA = f"/home/robot-lab/data/processed/{TASK}/{EP}"
OUT_DIR = "/home/robot-lab/cameron/yam_overlay"
os.makedirs(OUT_DIR, exist_ok=True)

for FRAME in FRAMES:
    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
    if not os.path.exists(pkl_path) or img_path is None:
        print(f"[frame {FRAME}] missing pkl or img, skip (pkl={os.path.exists(pkl_path)}, img={img_path})")
        continue
    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]  # right_arm_base origin in left_base/world

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

    def proj(p_world, label, color, radius=12):
        p_cam = (T_world2cam @ np.append(p_world, 1.0))[:3]
        if p_cam[2] < 0.01:
            print(f"  {label}: BEHIND CAM ({p_cam})"); return
        uv = K @ (p_cam / p_cam[2])
        u, v = uv[0], uv[1]
        oob = not (0 <= u < W and 0 <= v < H)
        tag = "OOB" if oob else "DRAWN"
        print(f"  {label}: world={p_world} uv=({u:.1f},{v:.1f}) {tag}")
        if not oob:
            cv2.drawMarker(img, (int(u), int(v)), color, cv2.MARKER_CROSS, 60, 4)
            cv2.circle(img, (int(u), int(v)), radius, color, 3)
            cv2.putText(img, label, (int(u)+14, int(v)-12),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)

    print(f"=== frame {FRAME} (task={TASK} ep={EP} cam={CAM}) ===")
    print(f"  l_grip={action[12]:.3f}  r_grip={action[25]:.3f}")
    print(f"  T_lfr translation (right_base in world): {r_base_world}")
    proj(l_pos,        "LEFT_EE",      (0, 255, 0))    # green
    proj(r_pos_world,  "RIGHT_EE",     (0, 0, 255))    # red
    proj(r_base_world, "RIGHT_BASE",   (255, 0, 255))  # magenta
    proj(np.array([0.0, 0.0, 0.0]), "WORLD_ORIG", (255, 255, 0))  # cyan-ish

    cv2.putText(img, f"{TASK}/{EP} f{FRAME} cam={CAM} v8 (raiden canonical)", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 2, cv2.LINE_AA)

    out_path = f"{OUT_DIR}/out_v8_{TASK}_{EP}_{CAM}_f{FRAME:04d}.png"
    cv2.imwrite(out_path, img)
    print(f"  saved {out_path}\n")