from __future__ import annotations

import numpy as np

from imu_lidar.contracts import ImuSeries
from rtk_imu.rtk_imu_multisource import R1bResult, UnifiedSession, solve_r2g


def _level_session(session_id: str) -> UnifiedSession:
    t = np.arange(0.0, 25.0, 0.01)
    return UnifiedSession(
        session_id=session_id,
        batch_id="test",
        imu=ImuSeries(
            t_s=t,
            gyro_rad_s=np.zeros((t.size, 3)),
            acc_m_s2=np.tile([0.0, 0.0, 9.80665], (t.size, 1)),
        ),
        imu_rpy_deg=np.zeros((t.size, 3)),
        imu_quaternion_wxyz=np.tile([1.0, 0.0, 0.0, 0.0], (t.size, 1)),
        imu_host_receive_utc_s=t,
        rtk_by_type={},
    )


def test_r2g_baseline_plus_level_gravity_completes_identity() -> None:
    r1b = R1bResult(
        baseline_axis_imu=np.array([1.0, 0.0, 0.0]),
        tilt_yz_deg=np.zeros(2),
        pair_count=100,
        residual_rms_deg=0.1,
        residual_p95_deg=0.2,
        covariance_deg2=np.eye(2) * 0.01,
        std_deg=np.ones(2) * 0.1,
        information_singular_values=np.ones(2),
        per_session_rms_deg={},
        gyro_bias_by_session_rad_s={},
        ok=True,
        notes=(),
    )
    sessions = [_level_session("a"), _level_session("b")]
    result = solve_r2g(sessions, r1b, level_static_session_ids={"a", "b"})
    assert result.R_RTK_IMU is not None
    assert np.allclose(result.R_RTK_IMU, np.eye(3), atol=1e-12)
    assert result.sample_count == 6
    assert result.session_count == 2