"""Unit tests for aruco_math.py — no ROS2, no live camera required."""

import math
import numpy as np
import pytest

from saltybot_aruco_detect.aruco_math import (
    rot_mat_to_quat,
    rvec_to_quat,
    tvec_distance,
    tvec_yaw_rad,
    MarkerPose,
)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _quat_norm(q):
    return math.sqrt(sum(x * x for x in q))


def _quat_rotate(q, v):
    """Rotate vector v by quaternion q (qx, qy, qz, qw)."""
    qx, qy, qz, qw = q
    # Quaternion product: q * (0,v) * q^-1
    # Using formula: v' = v + 2qw(q × v) + 2(q × (q × v))
    qvec = np.array([qx, qy, qz])
    t    = 2.0 * np.cross(qvec, v)
    return v + qw * t + np.cross(qvec, t)


# ── rot_mat_to_quat ───────────────────────────────────────────────────────────

def test_identity_rotation():
    R = np.eye(3)
    q = rot_mat_to_quat(R)
    assert abs(q[3] - 1.0) < 1e-9   # w = 1 for identity
    assert abs(q[0]) < 1e-9
    assert abs(q[1]) < 1e-9
    assert abs(q[2]) < 1e-9


def test_90deg_yaw():
    """90° rotation about Z axis."""
    angle = math.pi / 2
    R = np.array([
        [math.cos(angle), -math.sin(angle), 0],
        [math.sin(angle),  math.cos(angle), 0],
        [0,                0,               1],
    ])
    qx, qy, qz, qw = rot_mat_to_quat(R)
    # For 90° Z rotation: q = (0, 0, sin45°, cos45°)
    assert abs(qx) < 1e-6
    assert abs(qy) < 1e-6
    assert abs(qz - math.sin(angle / 2)) < 1e-6
    assert abs(qw - math.cos(angle / 2)) < 1e-6


def test_unit_norm():
    for yaw in [0, 0.5, 1.0, 2.0, math.pi]:
        R = np.array([
            [math.cos(yaw), -math.sin(yaw), 0],
            [math.sin(yaw),  math.cos(yaw), 0],
            [0, 0, 1],
        ])
        q = rot_mat_to_quat(R)
        assert abs(_quat_norm(q) - 1.0) < 1e-9, f"yaw={yaw}: norm={_quat_norm(q)}"


def test_roundtrip_rotation_matrix():
    """Rotating a vector with the matrix and quaternion should give same result."""
    angle = 1.23
    R = np.array([
        [math.cos(angle), -math.sin(angle), 0],
        [math.sin(angle),  math.cos(angle), 0],
        [0, 0, 1],
    ])
    q = rot_mat_to_quat(R)
    v = np.array([1.0, 0.0, 0.0])
    v_R = R @ v
    v_q = _quat_rotate(q, v)
    np.testing.assert_allclose(v_R, v_q, atol=1e-9)


def test_180deg_pitch():
    """180° rotation about Y axis."""
    R = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]], dtype=float)
    q = rot_mat_to_quat(R)
    assert abs(_quat_norm(q) - 1.0) < 1e-9
    # For 180° Y: q = (0, 1, 0, 0) or (0, -1, 0, 0)
    assert abs(abs(q[1]) - 1.0) < 1e-6


# ── rvec_to_quat ──────────────────────────────────────────────────────────────

def test_rvec_zero_is_identity():
    q = rvec_to_quat(np.array([0.0, 0.0, 0.0]))
    assert abs(_quat_norm(q) - 1.0) < 1e-6
    assert abs(q[3] - 1.0) < 1e-6   # w = 1


def test_rvec_unit_norm():
    for rvec in [
        [0.1, 0.0, 0.0],
        [0.0, 0.5, 0.0],
        [0.0, 0.0, 1.57],
        [0.3, 0.4, 0.5],
    ]:
        q = rvec_to_quat(np.array(rvec))
        assert abs(_quat_norm(q) - 1.0) < 1e-9, f"rvec={rvec}: norm={_quat_norm(q)}"


def test_rvec_z_rotation_matches_rot_mat():
    """rvec=[0,0,π/3] should match 60° Z rotation matrix quaternion."""
    angle = math.pi / 3
    rvec  = np.array([0.0, 0.0, angle])
    R     = np.array([
        [math.cos(angle), -math.sin(angle), 0],
        [math.sin(angle),  math.cos(angle), 0],
        [0, 0, 1],
    ])
    q_rvec = rvec_to_quat(rvec)
    q_R    = rot_mat_to_quat(R)
    # Allow sign flip (q and -q represent the same rotation)
    diff = min(
        abs(_quat_norm(np.array(q_rvec) - np.array(q_R))),
        abs(_quat_norm(np.array(q_rvec) + np.array(q_R))),
    )
    assert diff < 1e-6, f"rvec q={q_rvec}, mat q={q_R}"


# ── tvec_distance ─────────────────────────────────────────────────────────────

def test_tvec_distance_basic():
    assert abs(tvec_distance(np.array([3.0, 4.0, 0.0])) - 5.0) < 1e-9


def test_tvec_distance_zero():
    assert tvec_distance(np.array([0.0, 0.0, 0.0])) == 0.0


def test_tvec_distance_forward():
    assert abs(tvec_distance(np.array([0.0, 0.0, 2.5])) - 2.5) < 1e-9


def test_tvec_distance_3d():
    t = np.array([1.0, 2.0, 3.0])
    expected = math.sqrt(14.0)
    assert abs(tvec_distance(t) - expected) < 1e-9


# ── tvec_yaw_rad ──────────────────────────────────────────────────────────────

def test_yaw_zero_when_centred():
    """Marker directly in front: tx=0 → yaw=0."""
    assert abs(tvec_yaw_rad(np.array([0.0, 0.0, 2.0]))) < 1e-9


def test_yaw_positive_right():
    """Marker to the right (tx>0) → positive yaw."""
    assert tvec_yaw_rad(np.array([1.0, 0.0, 2.0])) > 0.0


def test_yaw_negative_left():
    """Marker to the left (tx<0) → negative yaw."""
    assert tvec_yaw_rad(np.array([-1.0, 0.0, 2.0])) < 0.0


def test_yaw_45deg():
    """tx=tz → 45°."""
    assert abs(tvec_yaw_rad(np.array([1.0, 0.0, 1.0])) - math.pi / 4) < 1e-9


# ── MarkerPose ────────────────────────────────────────────────────────────────

def _make_marker(mid=42, tx=0.0, ty=0.0, tz=2.0):
    rvec = np.array([0.0, 0.0, 0.1])
    tvec = np.array([tx, ty, tz])
    corners = np.zeros((4, 2))
    return MarkerPose(marker_id=mid, tvec=tvec, rvec=rvec, corners=corners)


def test_marker_distance_cached():
    m = _make_marker(tz=3.0)
    d1 = m.distance_m
    d2 = m.distance_m
    assert d1 == d2


def test_marker_lateral_and_forward():
    m = _make_marker(tx=0.5, tz=2.0)
    assert abs(m.lateral_m  - 0.5) < 1e-9
    assert abs(m.forward_m  - 2.0) < 1e-9


def test_marker_quat_unit_norm():
    m = _make_marker()
    q = m.quat
    assert abs(_quat_norm(q) - 1.0) < 1e-9


def test_marker_yaw_sign():
    m_right = _make_marker(tx=+0.3, tz=1.0)
    m_left  = _make_marker(tx=-0.3, tz=1.0)
    assert m_right.yaw_rad > 0
    assert m_left.yaw_rad  < 0
