"""Loop-closure self-test for both AR render layers (no cameras/servos needed):
for each layer, measure its board frame, render the board from a known angled pose,
localize that render, and check the recovered camera pose matches ground truth."""
import os
import sys

import numpy as np
import cv2
import mujoco

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from configs import CONFIGS                                  # noqa: E402
from render_core import Layer, look_at_mj, mat2quat, board_geom_pose  # noqa: E402
from calib_utils import cam_mj_to_cv, detect_boards, solve_pose       # noqa: E402
from aruco_boards import ARM_BASE, UMI_V2                    # noqa: E402

cases = [("scene", CONFIGS["arm"]["model"], ARM_BASE, "armboard"),
         ("umi", CONFIGS["gripper"]["model"], UMI_V2, "ar_back")]

for lbl, mp, board, mat in cases:
    L = Layer(lbl, mp, board, mat, lambda jn: jn)
    print(f"{lbl}: board '{board['name']}' measured, reproj {L.reproj:.2f}px")
    m, d, ext = L.model, L.data, L.extcam
    pos, xmat = board_geom_pose(m, d, mat)
    W = H = 900
    fovy = 50.0
    f = H / (2 * np.tan(np.radians(fovy) / 2))
    K = np.array([[f, 0, W / 2], [0, f, H / 2], [0, 0, 1]])
    done = False
    for sgn in (1.0, -1.0):
        eye = pos + xmat[:, 2] * sgn * 0.22 + xmat[:, 0] * 0.05
        p, Rmj = look_at_mj(eye, pos, up=xmat[:, 1])
        m.cam_pos[ext] = p
        m.cam_quat[ext] = mat2quat(Rmj)
        m.cam_fovy[ext] = fovy
        mujoco.mj_forward(m, d)
        Ttrue = cam_mj_to_cv(d.cam_xpos[ext], d.cam_xmat[ext])
        with mujoco.Renderer(m, H, W) as r:
            r.update_scene(d, camera=ext)
            img = r.render()
        dets = detect_boards(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), [board])
        if board["name"] in dets:
            T_B_cam, _ = solve_pose(*dets[board["name"]], K)
            Trec = L.T_board_w @ np.linalg.inv(T_B_cam)
            print(f"  loop-closure recover err: {np.linalg.norm(Trec[:3, 3] - Ttrue[:3, 3]) * 1000:.2f} mm")
            done = True
            break
    if not done:
        print("  recover: board not visible from test pose")
    L.close()