import cv2
import numpy as np

B2_CAMERA_CALIBRATION = load_npy_camera_calibration_file


def cam_to_world(cam_vector):
    assert cam_vector.shape[0] == 3

    tvec = B2_CAMERA_CALIBRATION['t']

    rotation_mtx = cv2.Rodrigues(B2_CAMERA_CALIBRATION['r'])[0]
    invert_rotation_mtx = np.linalg.inv(rotation_mtx)

    return invert_rotation_mtx @ (cam_vector - tvec)


def compute_zc(uv, z=0):
    assert uv.shape[0] == 3

    rotation_mtx = cv2.Rodrigues(B2_CAMERA_CALIBRATION['r'])[0]
    cross = np.cross(rotation_mtx[:, 0], rotation_mtx[:, 1]).reshape(3, 1)
    det = np.linalg.det(rotation_mtx)

    tvec = B2_CAMERA_CALIBRATION['t']

    intrinsic_mtx = B2_CAMERA_CALIBRATION['Intrinsic']
    invert_intrinsic_mtx = np.linalg.inv(intrinsic_mtx)

    return (np.dot(tvec.T, cross) + z * det) / np.dot(cross.T, invert_intrinsic_mtx @ uv)


def get_undistorted_points(coord):
    if coord.shape[0] == 3:
        coord = coord[:2, :]
    coord = coord.astype(np.float32)

    upt = cv2.undistortPoints(coord, B2_CAMERA_CALIBRATION['Intrinsic'], B2_CAMERA_CALIBRATION['dist_coeff'], None,
                              None, B2_CAMERA_CALIBRATION['Intrinsic'])
    return np.append(upt.reshape(-1, 2), np.ones([upt.shape[0], 1]), axis=1).T


def _2D_to_planar_world(uv, z=0):
    assert uv.shape[0] in [2, 3]

    uv = get_undistorted_points(uv)

    intrinsic_mtx = B2_CAMERA_CALIBRATION['Intrinsic']
    invert_intrinsic_mtx = np.linalg.inv(intrinsic_mtx)

    zc = compute_zc(uv, z)

    uvz = uv * zc

    cam_vector = invert_intrinsic_mtx @ uvz

    return cam_to_world(cam_vector)