import numpy as np
import cv2
import pytest
from aruco.pose_math import (
    rvec_tvec_to_matrix,
    matrix_to_rvec_tvec,
    invert_transform,
    compose_transforms,
    compute_reprojection_error,
)


def test_rvec_tvec_roundtrip():
    rvec = np.array([0.1, 0.2, 0.3])
    tvec = np.array([1.0, 2.0, 3.0])

    T = rvec_tvec_to_matrix(rvec, tvec)
    rvec_out, tvec_out = matrix_to_rvec_tvec(T)

    np.testing.assert_allclose(rvec, rvec_out, atol=1e-10)
    np.testing.assert_allclose(tvec, tvec_out, atol=1e-10)


def test_invert_transform():
    rvec = np.array([0.5, -0.2, 0.1])
    tvec = np.array([10.0, -5.0, 2.0])
    T = rvec_tvec_to_matrix(rvec, tvec)

    T_inv = invert_transform(T)
    I_test = T @ T_inv

    np.testing.assert_allclose(I_test, np.eye(4), atol=1e-10)


def test_compose_transforms():
    T1 = rvec_tvec_to_matrix(np.array([0.1, 0, 0]), np.array([1, 0, 0]))
    T2 = rvec_tvec_to_matrix(np.array([0, 0.2, 0]), np.array([0, 2, 0]))

    T_res = compose_transforms(T1, T2)
    T_expected = T1 @ T2

    np.testing.assert_allclose(T_res, T_expected, atol=1e-10)


def test_compute_reprojection_error_zero():
    # Setup camera
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
    dist = np.zeros(5)

    # Setup pose
    rvec = np.array([0.1, -0.2, 0.3])
    tvec = np.array([0.0, 0.0, 2.0])

    # Create object points
    obj_pts = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0]], dtype=np.float64)

    # Project them to get "perfect" image points
    img_pts, _ = cv2.projectPoints(obj_pts, rvec, tvec, K, dist)

    # Error should be near zero
    error = compute_reprojection_error(obj_pts, img_pts, rvec, tvec, K, dist)
    assert error < 1e-10


def test_compute_reprojection_error_nonzero():
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
    rvec = np.zeros(3, dtype=np.float64)
    tvec = np.array([0, 0, 1], dtype=np.float64)
    obj_pts = np.array([[0, 0, 0]], dtype=np.float64)

    # Projected point should be at (640, 360)
    # Let's provide an image point at (641, 360)
    img_pts = np.array([[641, 360]], dtype=np.float64)

    error = compute_reprojection_error(obj_pts, img_pts, rvec, tvec, K)
    assert abs(error - 1.0) < 1e-10