"""Tests for cached, session-balanced Phase-A comparison."""

from __future__ import annotations

import numpy as np

from imu_lidar.contracts import ImuSeries, MotionPair
from imu_lidar.geometry import so3_exp, so3_log
from imu_lidar.imu_preintegration import preintegrate_gyro
from imu_lidar.phase_a import (
    rehydrate_phase_a_pairs,
    select_decorrelated_phase_a_pairs,
    solve_phase_a_comparison,
)


def _phase_a_pair(
    session_id: str,
    index: int,
    r_true: np.ndarray,
    vector_deg: tuple[float, float, float],
    bias0: np.ndarray,
) -> MotionPair:
    r_b = so3_exp(np.deg2rad(np.asarray(vector_deg, dtype=float)))
    return MotionPair(
        session_id=session_id,
        i=index,
        j=index + 1,
        t_i_s=float(index),
        t_j_s=float(index + 1),
        R_A=r_true @ r_b @ r_true.T,
        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-5).tolist(),
            "gyro_bias0_rad_s": bias0.tolist(),
        },
    )


def test_phase_a_reports_three_variants_and_leave_one_session_out() -> None:
    r_true = so3_exp(np.deg2rad(np.array([3.0, -2.0, 25.0])))
    prior = so3_exp(np.deg2rad(np.array([0.0, 0.0, 0.2]))) @ r_true
    vectors = (
        (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),
        (9.0, -5.0, 6.0),
        (-6.0, 8.0, 11.0),
        (5.0, 7.0, -9.0),
    )
    biases = {
        "s0": np.array([0.001, -0.0005, 0.0002]),
        "s1": np.array([-0.0004, 0.0008, -0.0001]),
        "s2": np.array([0.0002, 0.0001, -0.0006]),
    }
    pairs: list[MotionPair] = []
    index = 0
    for sid, count in (("s0", 18), ("s1", 9), ("s2", 6)):
        for local_index in range(count):
            pairs.append(
                _phase_a_pair(
                    sid,
                    index,
                    r_true,
                    vectors[local_index % len(vectors)],
                    biases[sid],
                )
            )
            index += 1

    result = solve_phase_a_comparison(
        pairs,
        gyro_bias_rad_s_by_session=biases,
        rotation_prior=prior,
        rotation_prior_sigma_deg=15.0,
        yaw_std_max_deg=1.0,
        leave_one_out_yaw_range_max_deg=1.0,
        data_prior_difference_max_deg=1.0,
        decorrelation_block_s=0.0,
        max_nfev=80,
    )

    assert result.accepted
    assert result.strong_pair_counts_per_session == {
        "s0": 18,
        "s1": 9,
        "s2": 6,
    }
    assert len(result.leave_one_out) == 3
    assert result.marginal_observability.rank == 3
    assert result.leave_one_out_yaw_range_deg < 0.1
    for variant in (
        result.fixed_bg_data_only,
        result.session_bg_data_only,
        result.session_bg_with_rotation_prior,
    ):
        error_deg = np.degrees(
            np.linalg.norm(
                so3_log(r_true.T @ variant.R_IMU_lidar)
            )
        )
        assert error_deg < 0.1


def test_rehydrate_phase_a_pairs_recovers_jacobian_without_lidar() -> None:
    t_s = np.linspace(0.0, 1.0, 201)
    gyro = np.tile(np.array([0.12, -0.04, 0.2]), (t_s.size, 1))
    bias0 = np.array([0.01, -0.005, 0.002])
    imu = ImuSeries(
        t_s=t_s,
        gyro_rad_s=gyro,
        acc_m_s2=np.zeros((t_s.size, 3)),
    )
    preint = preintegrate_gyro(t_s, gyro, 0.1, 0.8, bias0)
    pair = MotionPair(
        session_id="s0",
        i=0,
        j=1,
        t_i_s=0.1,
        t_j_s=0.8,
        R_A=preint.delta_R,
        R_B=preint.delta_R,
        metadata={
            "t_i_imu_s": 0.1,
            "t_j_imu_s": 0.8,
            "gyro_bias0_rad_s": bias0.tolist(),
            "preint_sigma_rad": preint.sigma_rad,
        },
    )

    enriched, report = rehydrate_phase_a_pairs(
        [pair],
        imu_by_session={"s0": imu},
        bias0_by_session={"s0": bias0},
    )

    assert "J_bg" in enriched[0].metadata
    assert "cov" in enriched[0].metadata
    assert report["max_R_A_error_deg"] < 1e-8


def test_phase_a_time_blocks_do_not_count_overlapping_pairs_as_independent() -> None:
    r_true = so3_exp(np.deg2rad(np.array([1.0, -2.0, 20.0])))
    bias = np.zeros(3)
    pairs = [
        _phase_a_pair("s0", index, r_true, (5.0 + index, 2.0, 1.0), bias)
        for index in range(9)
    ]
    selected = select_decorrelated_phase_a_pairs(
        pairs,
        block_s=3.0,
        max_pairs_per_block=1,
    )
    assert len(selected) == 3
    assert all(pair in pairs for pair in selected)


def test_phase_a_planar_motion_is_partial_and_keeps_weak_direction_from_prior() -> None:
    r_true = so3_exp(np.deg2rad(np.array([4.0, -3.0, 31.0])))
    prior = so3_exp(np.deg2rad(np.array([0.2, -0.1, 0.4]))) @ r_true
    biases = {"s0": np.zeros(3), "s1": np.zeros(3)}
    pairs: list[MotionPair] = []
    for session_index, sid in enumerate(biases):
        for index in range(12):
            pairs.append(
                _phase_a_pair(
                    sid,
                    session_index * 100 + index,
                    r_true,
                    (0.0, 0.0, 8.0 + index),
                    biases[sid],
                )
            )
    result = solve_phase_a_comparison(
        pairs,
        gyro_bias_rad_s_by_session=biases,
        rotation_prior=prior,
        decorrelation_block_s=0.0,
        yaw_std_max_deg=0.5,
        run_leave_one_out=False,
        max_nfev=80,
    )
    assert not result.accepted
    assert result.partial_accepted
    assert result.solution_status == "phase_a_partial_accepted"
    assert result.marginal_observability.precision_rank == 2
    assert result.observable_subspace_with_prior is not None
    assert np.isinf(result.marginal_observability.direction_std_deg[0])