"""Pinhole calibration + board-pose utilities for the capture-relocalization pipeline."""
import numpy as np, cv2
from aruco_boards import DICT, ARM_BASE, UMI_TOP, UMI_BACK, board as mkboard

DETECTOR = cv2.aruco.ArucoDetector(DICT, cv2.aruco.DetectorParameters())
RX180 = np.diag([1.0, -1.0, -1.0])            # OpenCV cam frame = MuJoCo cam frame @ RX180


def detect_boards(gray, specs):
    """-> {name: (objpts Nx3, imgpts Nx2)} for each board with >=4 markers found."""
    corners, ids, _ = DETECTOR.detectMarkers(gray)
    out = {}
    if ids is None:
        return out
    for spec in specs:
        b = mkboard(spec)
        obj, img = b.matchImagePoints(corners, ids)
        if obj is not None and len(obj) >= 12:                # >=1 marker = 4 pts; want >=3 markers
            out[spec["name"]] = (obj.reshape(-1, 3), img.reshape(-1, 2))
    return out


def estimate_focal(views, size_wh, f_guess):
    """Zhang with principal point fixed at center, square pixels, zero distortion.
    views = list of (objpts, imgpts). Returns K, per-view (rvec, tvec), reproj err px."""
    W, H = size_wh
    K0 = np.array([[f_guess, 0, W / 2.0], [0, f_guess, H / 2.0], [0, 0, 1]], float)
    flags = (cv2.CALIB_USE_INTRINSIC_GUESS | cv2.CALIB_FIX_PRINCIPAL_POINT |
             cv2.CALIB_FIX_ASPECT_RATIO | cv2.CALIB_ZERO_TANGENT_DIST |
             cv2.CALIB_FIX_K1 | cv2.CALIB_FIX_K2 | cv2.CALIB_FIX_K3)
    obj = [v[0].astype(np.float32) for v in views]
    img = [v[1].astype(np.float32) for v in views]
    err, K, dist, rvecs, tvecs = cv2.calibrateCamera(obj, img, (W, H), K0, np.zeros(5), flags=flags)
    return K, list(zip(rvecs, tvecs)), err


def solve_pose(objpts, imgpts, K):
    ok, rvec, tvec = cv2.solvePnP(objpts.astype(np.float32), imgpts.astype(np.float32), K, None,
                                  flags=cv2.SOLVEPNP_ITERATIVE)
    proj, _ = cv2.projectPoints(objpts.astype(np.float32), rvec, tvec, K, None)
    err = float(np.linalg.norm(proj.reshape(-1, 2) - imgpts, axis=1).mean())
    return rt_to_T(rvec, tvec), err


def rt_to_T(rvec, tvec):
    T = np.eye(4)
    T[:3, :3] = cv2.Rodrigues(np.asarray(rvec, float))[0]
    T[:3, 3] = np.asarray(tvec, float).ravel()
    return T


def K_from_fovy(fovy_deg, W, H):
    f = H / (2.0 * np.tan(np.radians(fovy_deg) / 2.0))
    return np.array([[f, 0, W / 2.0], [0, f, H / 2.0], [0, 0, 1]], float)


def cam_mj_to_cv(pos, xmat):
    """MuJoCo camera world pose -> T_camCV_in_world (4x4)."""
    T = np.eye(4)
    T[:3, :3] = np.asarray(xmat, float).reshape(3, 3) @ RX180
    T[:3, 3] = np.asarray(pos, float)
    return T


def cam_cv_to_mj(T_camCV_in_world):
    """-> (pos, xyaxes string) for a MuJoCo <camera>."""
    R = T_camCV_in_world[:3, :3] @ RX180                       # back to MJ camera axes
    pos = T_camCV_in_world[:3, 3]
    x, y = R[:, 0], R[:, 1]
    return pos, " ".join(f"{v:.6f}" for v in list(x) + list(y))