"""Re-run ArUco detection at HD1080 with the new 3x3 board on the YAM."""
import os
os.environ.setdefault("MUJOCO_GL", "egl")

import sys, json, time
from pathlib import Path

import cv2
import numpy as np
import pyzed.sl as sl

sys.path.insert(0, "/home/robot-lab/cameron/raiden_fork/third_party/exo_redo")
from ExoConfigs.yam_exo import YAM_BASE_ONLY_CONFIG
from ExoConfigs.exoskeleton import link_to_aruco_transform
from exo_utils import do_est_aruco_pose, ARUCO_DICT

with open(os.path.expanduser("~/.config/raiden/camera.json")) as f:
    cam_cfg = json.load(f)
scene_serial = next(int(e["serial"]) for e in cam_cfg.values() if e.get("role") == "scene")
print(f"Scene cam serial: {scene_serial}\n")

cfg = YAM_BASE_ONLY_CONFIG
link_cfg = cfg.links["larger_coarse_board"]
aruco_board = cfg.aruco_board_objects["larger_coarse_board"]
board_length = link_cfg.board_length
T_link_to_aruco = link_to_aruco_transform(link_cfg)
T_aruco_to_link = np.linalg.inv(T_link_to_aruco)

OUT_DIR = Path("/home/robot-lab/cameron/yam_overlay/aruco_probe")
OUT_DIR.mkdir(parents=True, exist_ok=True)

cam = sl.Camera()
init = sl.InitParameters()
init.set_from_serial_number(scene_serial)
init.camera_resolution = sl.RESOLUTION.HD1080
init.camera_fps = 15
init.depth_mode = sl.DEPTH_MODE.NONE
init.coordinate_units = sl.UNIT.METER

status = cam.open(init)
if status != sl.ERROR_CODE.SUCCESS:
    sys.exit(f"open failed: {status}")

info = cam.get_camera_information()
cal = info.camera_configuration.calibration_parameters.left_cam
W, H = info.camera_configuration.resolution.width, info.camera_configuration.resolution.height
K = np.array([[cal.fx, 0, cal.cx], [0, cal.fy, cal.cy], [0, 0, 1]], dtype=np.float64)
dist = np.array([cal.disto[0], cal.disto[1], cal.disto[2], cal.disto[3], cal.disto[4]], dtype=np.float64)
print(f"HD1080: {W}x{H},  K diag=({cal.fx:.1f}, {cal.fy:.1f}),  cxcy=({cal.cx:.1f}, {cal.cy:.1f})\n")
print(f"Board: 3x3, IDs 217-225, board_length={board_length*1000:.1f} mm")
print(f"link→aruco offset: {link_cfg.aruco_offset_pos} mm\n")

# Warm up + grab
image = sl.Mat()
rt = sl.RuntimeParameters()
bgr = None
for i in range(20):
    if cam.grab(rt) != sl.ERROR_CODE.SUCCESS:
        time.sleep(0.05)
        continue
    cam.retrieve_image(image, sl.VIEW.LEFT)
    npm = image.get_data()
    bgr = cv2.cvtColor(npm, cv2.COLOR_BGRA2BGR)

# Save raw
raw_path = OUT_DIR / "frame_3x3_HD1080.png"
cv2.imwrite(str(raw_path), bgr)
print(f"  raw frame → {raw_path}")

# Low-level marker detection
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
corners, ids, rejected = cv2.aruco.detectMarkers(
    gray, ARUCO_DICT, parameters=cv2.aruco.DetectorParameters()
)
n_markers = 0 if ids is None else len(ids)
n_rejected = 0 if rejected is None else len(rejected)
detected_ids = sorted(ids.flatten().tolist()) if ids is not None else []
print(f"  cv2.aruco.detectMarkers → {n_markers} markers, {n_rejected} rejected")
print(f"  detected IDs: {detected_ids}")

# Pose estimation
try:
    result = do_est_aruco_pose(bgr, ARUCO_DICT, aruco_board, board_length,
                               cameraMatrix=K, distCoeffs=dist)
except Exception as e:
    result = -1
    print(f"  do_est_aruco_pose error: {e}")

if result == -1:
    print(f"  POSE: FAILED")
    vis = bgr.copy()
    cv2.putText(vis, f"3x3 HD1080: pose failed (markers={n_markers})",
                (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)
else:
    T_aruco_in_cam = result["est_aruco_pose"]
    T_link_in_cam = T_aruco_in_cam @ T_aruco_to_link
    T_cam_in_base = np.linalg.inv(T_link_in_cam)
    t = T_cam_in_base[:3, 3]
    obj_cam, img_pts = result["obj_img_pts"]
    rvec, tvec = result["rtvec"]
    R_mat, _ = cv2.Rodrigues(rvec)
    obj_pts_world = (R_mat.T @ (obj_cam.T - tvec.reshape(3, 1))).T
    proj, _ = cv2.projectPoints(obj_pts_world, rvec, tvec, K, dist)
    residual = np.linalg.norm(proj.reshape(-1, 2) - img_pts.reshape(-1, 2), axis=1).mean()
    print(f"\n  POSE OK")
    print(f"    cam_in_base xyz = ({t[0]:+.3f}, {t[1]:+.3f}, {t[2]:+.3f}) m")
    print(f"    mean reproj resid = {residual:.2f} px")
    vis = result["pose_vis"]

annot_path = OUT_DIR / "frame_3x3_HD1080_annot.png"
cv2.imwrite(str(annot_path), vis)
print(f"\n  annotated frame → {annot_path}")

cam.close()
print("\nBrowse: /browse/yam_remote/cameron/yam_overlay/aruco_probe/")