"""ArUco detection at HD1080 with full exo+robot mesh overlay (like exo_cam.py)."""
import os
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")

import sys, json, time
from pathlib import Path

import cv2
import mujoco
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,
    position_exoskeleton_meshes, get_link_poses_from_robot, render_from_camera_pose,
)


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}")

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)
print(f"Board: 3x3, IDs 217-225, side={board_length*1000:.1f}mm")

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

# Build the mujoco model from the exo config (robot+exo+aruco plane mocap bodies).
# Inject balanceinertia so the degenerate ArUco-plane inertia tensor doesn't
# fail MuJoCo's triangle-inequality validation (these meshes are mocap-driven
# visuals only — no collision or dynamics).
# Also chdir to exo_redo root so the XML's relative texture/mesh paths
# (e.g. ``robot_models/board_imgs/larger_coarse_board.png``) resolve.
print("Building mujoco model from YAM_BASE_ONLY_CONFIG...")
os.chdir("/home/robot-lab/cameron/raiden_fork/third_party/exo_redo")
xml_str = cfg.xml.replace(
    '<compiler angle="radian"',
    '<compiler angle="radian" balanceinertia="true"',
)
import re
xml_str = re.sub(r'<include file="[^"]*background[^"]*"\s*/>', '', xml_str)
model = mujoco.MjModel.from_xml_string(xml_str)
data = mujoco.MjData(model)
mujoco.mj_forward(model, data)

# Position the base "arm" body at world origin (it's the base; no joints needed).
# For YAM with only a base body, mj_forward already puts the robot at default qpos.
# Then place exo/aruco-plane mocaps via link_poses from the (zero) robot state.
link_poses = get_link_poses_from_robot(cfg, model, data)
position_exoskeleton_meshes(cfg, model, data, link_poses)

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

if cam.open(init) != sl.ERROR_CODE.SUCCESS:
    sys.exit("ZED open failed")

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})\n")

# Warm up
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)
    bgr = cv2.cvtColor(image.get_data(), cv2.COLOR_BGRA2BGR)

raw_path = OUT_DIR / "frame_3x3_HD1080_overlay_raw.png"
cv2.imwrite(str(raw_path), bgr)

# Detect
result = do_est_aruco_pose(bgr, ARUCO_DICT, aruco_board, board_length, cameraMatrix=None)

if result == -1:
    print("POSE FAILED")
    cam.close()
    sys.exit(1)

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]

# Reproj residual
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"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\n")

# Render the full robot+exo+aruco-plane scene from the recovered camera pose
print("Rendering robot+exo+aruco mesh overlay...")
K_solved = result["cameraMatrix"]; rendered = render_from_camera_pose(model, data, T_link_in_cam, K_solved, H, W)  # RGB uint8

# Composite: alpha-blend the rendered scene over the RGB frame.
# Rendered pixels that are dark (background) get less weight.
rgb_frame = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB).astype(np.float32)
rendered_f = rendered.astype(np.float32)
# Background-mask: pixels close to black in render = background; full weight on RGB there
bg_mask = (rendered_f.sum(axis=-1) < 30).astype(np.float32)[..., None]
# Composite: where render has content, blend 0.55 render + 0.45 RGB; else 100% RGB
composite = (1 - bg_mask) * (0.7 * rendered_f + 0.3 * rgb_frame) + bg_mask * rgb_frame
composite = composite.clip(0, 255).astype(np.uint8)

# Convert to BGR for cv2 + draw text
composite_bgr = cv2.cvtColor(composite, cv2.COLOR_RGB2BGR)
cv2.putText(composite_bgr,
            f"cam@base=({t[0]:+.3f},{t[1]:+.3f},{t[2]:+.3f})  resid={residual:.2f}px",
            (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2, cv2.LINE_AA)

# Also keep the marker-annotated frame
marker_annot = result["pose_vis"]

# Three-panel view: raw | markers+axes | mesh-overlay
panel = np.concatenate([bgr, marker_annot, composite_bgr], axis=1)
out_path = OUT_DIR / "frame_3x3_HD1080_overlay_panel.png"
cv2.imwrite(str(out_path), panel)
print(f"  panel (raw | aruco | mesh-overlay) → {out_path}")

# Also save just the mesh overlay alone
overlay_path = OUT_DIR / "frame_3x3_HD1080_overlay_only.png"
cv2.imwrite(str(overlay_path), composite_bgr)
print(f"  overlay alone               → {overlay_path}")

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