"""Regression tests for calibration quality, continuity, and observability gates."""

from __future__ import annotations

import numpy as np

from imu_lidar.contracts import ImuSeries, LidarFrame, MotionPair
from imu_lidar.geometry import make_transform, so3_exp, so3_log
from imu_lidar.joint_optimizer import solve_joint_extrinsic
from imu_lidar.motion_pairs import build_motion_pairs
from imu_lidar.observability import analyze_observability
from imu_lidar.registration import RegistrationResult
from imu_lidar.rotation_handeye import solve_rotation_handeye


def _motion_pair(index: int, rotation_vector: np.ndarray) -> MotionPair:
    rotation = so3_exp(np.asarray(rotation_vector, dtype=float))
    return MotionPair(
        session_id="synthetic",
        i=index,
        j=index + 1,
        t_i_s=float(index),
        t_j_s=float(index + 1),
        R_A=rotation,
        R_B=rotation,
        t_A_m=np.zeros(3),
        t_B_m=np.array([0.1, -0.03, 0.0]),
        fitness=0.9,
        metadata={
            "J_bg": (-np.eye(3)).tolist(),
            "cov": (np.eye(3) * 1e-4).tolist(),
            "gyro_bias0_rad_s": [0.0, 0.0, 0.0],
        },
    )


def _frame(frame_id: str, mid_s: float) -> LidarFrame:
    return LidarFrame(
        frame_id=frame_id,
        t_start_s=mid_s - 0.01,
        t_end_s=mid_s + 0.01,
        points_xyz=np.zeros((64, 3)),
    )


def _registration(*, fitness: float = 0.9) -> RegistrationResult:
    rotation = so3_exp(np.deg2rad(np.array([0.0, 0.0, 10.0])))
    return RegistrationResult(
        transform=make_transform(np.array([0.4, 0.0, 0.0]), rotation),
        fitness=fitness,
        rotation_deg=10.0,
        translation_m=0.4,
        backend="test",
        ok=True,
    )


def test_planar_yaw_is_not_full_rotation_or_translation_observable():
    pairs = [
        _motion_pair(i, np.deg2rad(np.array([0.0, 0.0, angle_deg])))
        for i, angle_deg in enumerate((5.0, 8.0, 12.0, 17.0, 23.0, 31.0))
    ]

    report = analyze_observability(pairs, np.eye(3))

    assert not report.rotation_observable
    assert not report.translation_observable


def test_multi_axis_motion_is_rotation_and_translation_observable():
    vectors_deg = (
        (12.0, 0.0, 0.0),
        (0.0, 15.0, 0.0),
        (0.0, 0.0, 18.0),
        (10.0, 8.0, 0.0),
        (0.0, 11.0, 9.0),
        (7.0, 0.0, 13.0),
    )
    pairs = [
        _motion_pair(i, np.deg2rad(np.asarray(vector_deg)))
        for i, vector_deg in enumerate(vectors_deg)
    ]

    report = analyze_observability(pairs, np.eye(3))

    assert report.rotation_observable
    assert report.translation_observable


def test_translation_prior_is_reported_but_not_accepted_when_unobservable():
    pairs = [
        _motion_pair(i, np.deg2rad(np.array([0.0, 0.0, angle_deg])))
        for i, angle_deg in enumerate((5.0, 8.0, 12.0, 17.0, 23.0, 31.0))
    ]
    prior = np.array([0.3, -0.2, 0.5])

    result = solve_joint_extrinsic(
        pairs,
        np.eye(3),
        force_rotation_only=False,
        enable_phase_c=False,
        t_prior_m=prior,
    )

    assert not result.translation_accepted
    np.testing.assert_allclose(result.T_IMU_lidar[:3, 3], prior)
    assert any("prior only" in note for note in result.notes)


def test_handeye_rejects_a_small_fraction_of_gross_rotation_outliers():
    rng = np.random.default_rng(7)
    r_true = so3_exp(np.deg2rad(np.array([2.0, -3.0, 20.0])))
    pairs: list[MotionPair] = []
    for index in range(100):
        axis = rng.normal(size=3)
        axis /= np.linalg.norm(axis)
        r_b = so3_exp(axis * np.deg2rad(rng.uniform(8.0, 30.0)))
        r_a = r_true @ r_b @ r_true.T
        if index == 0:
            r_a = so3_exp(np.deg2rad(np.array([18.0, 0.0, 0.0]))) @ r_a
        pairs.append(
            MotionPair(
                session_id="outlier",
                i=index,
                j=index + 1,
                t_i_s=float(index),
                t_j_s=float(index + 1),
                R_A=r_a,
                R_B=r_b,
            )
        )

    result = solve_rotation_handeye(pairs)

    assert not result.ok
    assert result.outlier_fraction_gt_5deg > 0.005


def test_motion_pairs_reject_low_fitness(monkeypatch):
    monkeypatch.setattr(
        "imu_lidar.motion_pairs.register_lidar_pair",
        lambda *_args, **_kwargs: _registration(fitness=0.3),
    )
    imu = ImuSeries(
        t_s=np.linspace(0.0, 1.2, 121),
        gyro_rad_s=np.zeros((121, 3)),
        acc_m_s2=np.zeros((121, 3)),
    )

    result = build_motion_pairs(
        session_id="fitness",
        keyframes=[_frame("0", 0.1), _frame("1", 1.1)],
        keyframe_indices=[0, 1],
        imu=imu,
        delta_t_s=0.0,
        min_registration_fitness=0.5,
    )

    assert not result.pairs
    assert any("fitness<0.50: 1" in note for note in result.notes)


def test_motion_pairs_reject_imu_and_lidar_discontinuities(monkeypatch):
    monkeypatch.setattr(
        "imu_lidar.motion_pairs.register_lidar_pair",
        lambda *_args, **_kwargs: _registration(),
    )
    imu_with_gap = ImuSeries(
        t_s=np.array([0.0, 0.1, 0.2, 0.3, 0.4, 0.8, 0.9, 1.0, 1.1, 1.2]),
        gyro_rad_s=np.zeros((10, 3)),
        acc_m_s2=np.zeros((10, 3)),
    )
    imu_result = build_motion_pairs(
        session_id="imu-gap",
        keyframes=[_frame("0", 0.1), _frame("1", 1.1)],
        keyframe_indices=[0, 1],
        imu=imu_with_gap,
        delta_t_s=0.0,
        max_imu_gap_s=0.2,
    )

    assert not imu_result.pairs
    assert any("IMU gap>0.200s: 1" in note for note in imu_result.notes)

    continuous_imu = ImuSeries(
        t_s=np.linspace(0.0, 2.2, 221),
        gyro_rad_s=np.zeros((221, 3)),
        acc_m_s2=np.zeros((221, 3)),
    )
    lidar_result = build_motion_pairs(
        session_id="lidar-gap",
        keyframes=[_frame("0", 0.1), _frame("2", 2.1)],
        keyframe_indices=[0, 2],
        imu=continuous_imu,
        delta_t_s=0.0,
        all_frame_times_s=np.array([0.1, 0.2, 2.1]),
        max_lidar_gap_s=0.5,
    )

    assert not lidar_result.pairs
    assert any("LiDAR gap>0.500s: 1" in note for note in lidar_result.notes)

def test_phase_a_keeps_session_bias_linearization_points_independent():
    r_true = so3_exp(np.deg2rad(np.array([2.0, -3.0, 20.0])))
    bias0_by_session = {
        "s0": np.array([0.010, -0.004, 0.002]),
        "s1": np.array([-0.006, 0.008, -0.003]),
    }
    vectors_deg = (
        (12.0, 0.0, 0.0),
        (0.0, 15.0, 0.0),
        (0.0, 0.0, 18.0),
        (10.0, 8.0, 0.0),
        (0.0, 11.0, 9.0),
        (7.0, 0.0, 13.0),
    )
    pairs: list[MotionPair] = []
    for session_index, (session_id, bias0) in enumerate(bias0_by_session.items()):
        for pair_index, vector_deg in enumerate(vectors_deg):
            r_b = so3_exp(np.deg2rad(np.asarray(vector_deg)))
            r_a = r_true @ r_b @ r_true.T
            index = session_index * 100 + pair_index
            pairs.append(
                MotionPair(
                    session_id=session_id,
                    i=index,
                    j=index + 1,
                    t_i_s=float(pair_index),
                    t_j_s=float(pair_index + 1),
                    R_A=r_a,
                    R_B=r_b,
                    t_A_m=np.zeros(3),
                    t_B_m=np.zeros(3),
                    metadata={
                        "J_bg": np.eye(3).tolist(),
                        "cov": (np.eye(3) * 1e-4).tolist(),
                        "gyro_bias0_rad_s": bias0.tolist(),
                    },
                )
            )

    result = solve_joint_extrinsic(
        pairs,
        r_true,
        force_rotation_only=True,
        gyro_bias_rad_s_by_session=bias0_by_session,
    )

    assert result.phase_a_accepted
    assert set(result.phase_a_comparison["variants"]) == {
        "A0_fixed_bg_data_only",
        "A1_session_bg_data_only",
        "A2_session_bg_with_rotation_prior",
    }
    assert set(result.gyro_bias_rad_s_per_session) == {"s0", "s1"}
    for session_id, bias0 in bias0_by_session.items():
        np.testing.assert_allclose(
            result.gyro_bias_rad_s_per_session[session_id], bias0, atol=1e-8
        )
    assert np.linalg.norm(so3_log(r_true.T @ result.T_IMU_lidar[:3, :3])) < 1e-8