import numpy as np
from scipy.spatial.transform import Rotation as R
from vid2rig.rotutil import aa_to_quat, fix_hemisphere, slerp_series

def test_aa_to_quat_matches_scipy():
    aa = np.array([[0, 0, np.pi / 2], [0.3, -0.2, 0.1]])
    assert np.allclose(np.abs(aa_to_quat(aa)), np.abs(R.from_rotvec(aa).as_quat()))

def test_fix_hemisphere_flips_sign():
    q = R.from_rotvec([0, 0, 0.1]).as_quat()
    series = np.stack([q, -q, q])           # artificially flipped middle frame
    fixed = fix_hemisphere(series)
    dots = np.sum(fixed[1:] * fixed[:-1], axis=1)
    assert np.all(dots >= -1e-6)

def test_slerp_midpoint():
    q0 = R.from_rotvec([0, 0, 0.0]).as_quat()
    q1 = R.from_rotvec([0, 0, np.pi / 2]).as_quat()
    out = slerp_series(np.stack([q0, q1]), np.array([0.0, 1.0]), np.array([0.5]))
    expected = R.from_rotvec([0, 0, np.pi / 4]).as_quat()
    assert np.allclose(np.abs(out[0]), np.abs(expected), atol=1e-6)