"""Tests for motion-pair cache IO."""

from __future__ import annotations

from pathlib import Path

import numpy as np

from imu_lidar.contracts import MotionPair
from imu_lidar.motion_pairs_io import (
    build_motion_pairs_payload,
    load_motion_pairs,
    pair_from_dict,
    pair_to_dict,
    pairs_for_session,
    save_motion_pairs,
)


def test_pair_roundtrip(tmp_path: Path) -> None:
    pair = MotionPair(
        session_id="s0",
        i=1,
        j=4,
        t_i_s=1.0,
        t_j_s=2.5,
        R_A=np.eye(3),
        R_B=np.eye(3),
        t_A_m=np.array([0.1, 0.0, 0.0]),
        t_B_m=np.array([0.1, 0.0, 0.0]),
        fitness=0.8,
        metadata={
            "weight": 12.0,
            "cov": (np.eye(3) * 1e-4).tolist(),
            "J_bg": (-np.eye(3)).tolist(),
            "cov9": [[0.0] * 9] * 9,
            "backend": "test",
            "gyro_bias0_rad_s": [0.01, -0.02, 0.03],
            "accel_bias0_m_s2": [0.1, 0.2, -0.1],
            "time_offset_s": 0.004,
            "keyframe_span": 3,
            "is_consecutive": False,
        },
    )
    encoded = pair_to_dict(pair)
    assert "cov9" not in encoded["metadata"]
    assert "cov" in encoded["metadata"]
    assert "J_bg" in encoded["metadata"]
    assert encoded["metadata"]["weight"] == 12.0
    restored = pair_from_dict(encoded)
    assert restored.i == 1 and restored.j == 4
    np.testing.assert_allclose(restored.t_A_m, [0.1, 0.0, 0.0])
    np.testing.assert_allclose(restored.metadata["gyro_bias0_rad_s"], [0.01, -0.02, 0.03])
    assert restored.metadata["keyframe_span"] == 3

    payload = build_motion_pairs_payload(
        prepared_sessions=[
            {
                "session_id": "s0",
                "time_offset_s": 0.0,
                "gyro_bias_rad_s": np.zeros(3),
                "pairs": (pair,),
            }
        ]
    )
    path = save_motion_pairs(tmp_path / "motion_pairs.json", payload)
    loaded = load_motion_pairs(path)
    assert loaded["schema_version"] == 2
    pairs = pairs_for_session(loaded, "s0")
    assert len(pairs) == 1
    assert pairs[0].session_id == "s0"

    payload["schema_version"] = 1
    legacy_path = save_motion_pairs(
        tmp_path / "motion_pairs_v1.json", payload
    )
    legacy = load_motion_pairs(legacy_path)
    assert legacy["schema_version"] == 1