import numpy as np
import cv2


def rvec_tvec_to_matrix(rvec, tvec):
    rvec = np.asarray(rvec).flatten()
    tvec = np.asarray(tvec).flatten()
    R, _ = cv2.Rodrigues(rvec)
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = tvec
    return T


def matrix_to_rvec_tvec(T):
    R = T[:3, :3]
    tvec = T[:3, 3]
    rvec, _ = cv2.Rodrigues(R)
    return rvec.flatten(), tvec.flatten()


def invert_transform(T):
    R = T[:3, :3]
    t = T[:3, 3]
    T_inv = np.eye(4)
    T_inv[:3, :3] = R.T
    T_inv[:3, 3] = -R.T @ t
    return T_inv


def compose_transforms(T1, T2):
    return T1 @ T2


def compute_reprojection_error(obj_pts, img_pts, rvec, tvec, K, dist=None):
    projected_pts, _ = cv2.projectPoints(obj_pts, rvec, tvec, K, dist)
    projected_pts = projected_pts.squeeze()
    img_pts = img_pts.squeeze()
    if projected_pts.ndim == 1:
        projected_pts = projected_pts.reshape(1, -1)
    if img_pts.ndim == 1:
        img_pts = img_pts.reshape(1, -1)
    error = np.linalg.norm(img_pts - projected_pts, axis=1)
    return float(np.mean(error))