import numpy as np
from vid2rig.motion import MotionClip

def _toy_clip():
    F, J = 3, 2
    return MotionClip(
        fps=30.0, num_frames=F, source="smpl",
        joint_names=["pelvis", "spine"],
        local_rotations=np.tile(np.array([0, 0, 0, 1.0]), (F, J, 1)),
        root_translation=np.arange(F * 3, dtype=float).reshape(F, 3),
        root_rotation=np.tile(np.array([0, 0, 0, 1.0]), (F, 1)),
        rest_offsets=np.array([[0, 0, 0], [0, 0.2, 0]], dtype=float),
        parents=[-1, 0],
        source_height_m=1.7,
    )

def test_roundtrip(tmp_path):
    c = _toy_clip()
    p = tmp_path / "clip.npz"
    c.save(p)
    d = MotionClip.load(p)
    assert d.num_frames == 3 and d.parents == [-1, 0] and d.source == "smpl"
    assert np.allclose(d.root_translation, c.root_translation)
    assert d.joint_names == ["pelvis", "spine"]

def test_root_rotation_mirrors_joint0():
    c = _toy_clip()
    assert np.allclose(c.local_rotations[:, 0], c.root_rotation)