"""Automated tests for imu_lidar (synthetic data).

See ``tests/README.md`` for:
- what each pytest covers;
- offline S2 host-time experiments (not run in default pytest) and recorded outcomes.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

from imu_lidar.contracts import CalibrationMode, CalibrationRequest, MotionPair, SessionInput
from imu_lidar.geometry import so3_exp
from imu_lidar.pipeline import run_calibration
from imu_lidar.rotation_handeye import solve_rotation_handeye
from imu_lidar.time_offset import estimate_time_offset
from imu_lidar.imu_io import load_imu_samples
from imu_lidar.lidar_io import load_lidar_frames
from tools.generate_synthetic_session import generate_synthetic_session


def test_rotation_handeye_recovers_yaw():
    r_true = so3_exp(np.deg2rad(np.array([1.0, -2.0, 30.0])))
    pairs = []
    rng = np.random.default_rng(1)
    for _ in range(20):
        axis = rng.normal(size=3)
        axis /= np.linalg.norm(axis)
        angle = np.deg2rad(rng.uniform(8.0, 35.0))
        r_b = so3_exp(axis * angle)
        r_a = r_true @ r_b @ r_true.T
        pairs.append(
            MotionPair(
                session_id="s",
                i=0,
                j=1,
                t_i_s=0.0,
                t_j_s=1.0,
                R_A=r_a,
                R_B=r_b,
            )
        )
    result = solve_rotation_handeye(pairs)
    assert result.ok
    err = np.linalg.norm(_log(r_true.T @ result.R_IMU_lidar))
    assert np.degrees(err) < 1.0


def _log(rotation: np.ndarray) -> np.ndarray:
    from imu_lidar.geometry import so3_log

    return so3_log(rotation)


def test_synthetic_pipeline_rotation_and_time_offset(tmp_path: Path):
    meta = generate_synthetic_session(tmp_path, delta_t_s=0.17, yaw_extrinsic_deg=25.0)
    config = Path(__file__).resolve().parents[1] / "config" / "vehicle_installation.template.yaml"
    out = tmp_path / "out"
    request = CalibrationRequest(
        vehicle_config=config,
        sessions=(
            SessionInput(
                session_id="synth",
                imu_source=tmp_path / "imu.csv",
                lidar_source=tmp_path / "lidar",
            ),
        ),
        requested_mode=CalibrationMode.ROTATION_ONLY,
        output_directory=out,
        max_iterations=1,
        time_offset_search_s=0.5,
        min_pair_rotation_deg=2.0,
        min_pair_translation_m=0.05,
    )
    result = run_calibration(request)
    assert result.status.value == "rotation_only_accepted"
    assert result.time_offset_s is not None
    assert abs(result.time_offset_s - meta["delta_t_s"]) < 0.05
    assert result.T_IMU_lidar is not None
    # End-to-end uses approximate ICP; allow moderate absolute error but require consistency.
    r_true = so3_exp(np.deg2rad(np.array([2.0, -1.5, meta["yaw_extrinsic_deg"]])))
    err_deg = np.degrees(np.linalg.norm(_log(r_true.T @ result.T_IMU_lidar[:3, :3])))
    assert err_deg < 15.0
    session0 = result.details["sessions"][0]
    assert session0["handeye"]["residual_rms_deg"] < 5.0


def test_time_offset_on_synthetic(tmp_path: Path):
    meta = generate_synthetic_session(tmp_path, delta_t_s=0.21, yaw_extrinsic_deg=15.0)
    imu = load_imu_samples(tmp_path / "imu.csv")
    frames = load_lidar_frames(tmp_path / "lidar")
    offset = estimate_time_offset(imu, frames, search_s=0.5)
    assert offset.ok
    assert abs(offset.delta_t_s - meta["delta_t_s"]) < 0.05


def test_preintegration_bias_jacobian_matches_finite_difference():
    from imu_lidar.imu_preintegration import apply_bias_jacobian_correction, preintegrate_gyro
    from imu_lidar.geometry import so3_log

    rng = np.random.default_rng(0)
    t = np.linspace(0.0, 1.0, 200)
    gyro = rng.normal(scale=0.2, size=(t.size, 3))
    bias0 = np.array([0.01, -0.02, 0.005])
    preint = preintegrate_gyro(t, gyro, 0.1, 0.7, bias0)
    db = np.array([1e-3, -2e-3, 5e-4])
    approx = apply_bias_jacobian_correction(preint.delta_R, preint.J_bg, db)
    exact = preintegrate_gyro(t, gyro, 0.1, 0.7, bias0 + db).delta_R
    err = np.linalg.norm(so3_log(approx.T @ exact))
    assert err < 2e-3


def test_imu_preintegration_recovers_constant_accel_translation():
    from imu_lidar.imu_preintegration import preintegrate_imu
    from imu_lidar.geometry import so3_log

    # Constant body accel (no gravity in preint body increments), zero gyro.
    dt = 0.01
    t = np.arange(0.0, 1.0 + 1e-9, dt)
    gyro = np.zeros((t.size, 3))
    acc = np.tile(np.array([0.5, -0.2, 0.1]), (t.size, 1))
    preint = preintegrate_imu(t, gyro, acc, 0.0, 1.0, np.zeros(3), np.zeros(3))
    assert np.linalg.norm(so3_log(preint.delta_R)) < 1e-9
    # Δv ≈ a Δt, Δp ≈ 0.5 a Δt²
    assert np.linalg.norm(preint.delta_v - acc[0] * 1.0) < 5e-3
    assert np.linalg.norm(preint.delta_p - 0.5 * acc[0] * 1.0) < 1e-2
    assert preint.cov.shape == (9, 9)
    assert preint.J_bg.shape == (9, 3) and preint.J_ba.shape == (9, 3)


def test_imu_preintegration_bias_jacobian_finite_difference():
    from imu_lidar.imu_preintegration import apply_bias_correction_imu, preintegrate_imu

    rng = np.random.default_rng(2)
    t = np.linspace(0.0, 0.8, 160)
    gyro = rng.normal(scale=0.15, size=(t.size, 3))
    acc = rng.normal(scale=0.5, size=(t.size, 3)) + np.array([0.0, 0.0, 9.8])
    bg0 = np.array([0.01, -0.01, 0.0])
    ba0 = np.array([0.02, 0.0, -0.01])
    base = preintegrate_imu(t, gyro, acc, 0.05, 0.55, bg0, ba0)
    dbg = np.array([5e-4, -3e-4, 2e-4])
    dba = np.array([1e-3, -5e-4, 0.0])
    r_a, v_a, p_a = apply_bias_correction_imu(base, dbg, dba)
    exact = preintegrate_imu(t, gyro, acc, 0.05, 0.55, bg0 + dbg, ba0 + dba)
    from imu_lidar.geometry import so3_log

    assert np.linalg.norm(so3_log(r_a.T @ exact.delta_R)) < 5e-3
    assert np.linalg.norm(v_a - exact.delta_v) < 3e-2
    assert np.linalg.norm(p_a - exact.delta_p) < 2e-2


def test_synthetic_pipeline_full_se3_smoke(tmp_path: Path):
    generate_synthetic_session(tmp_path, delta_t_s=0.12, yaw_extrinsic_deg=18.0)
    config = Path(__file__).resolve().parents[1] / "config" / "vehicle_installation.template.yaml"
    out = tmp_path / "out_se3"
    request = CalibrationRequest(
        vehicle_config=config,
        sessions=(
            SessionInput(
                session_id="synth",
                imu_source=tmp_path / "imu.csv",
                lidar_source=tmp_path / "lidar",
            ),
        ),
        requested_mode=CalibrationMode.FULL_SE3,
        output_directory=out,
        max_iterations=1,
        time_offset_search_s=0.5,
        min_pair_rotation_deg=2.0,
        min_pair_translation_m=0.05,
    )
    result = run_calibration(request)
    assert result.status.value in {
        "full_se3_accepted",
        "full_se3_rejected_due_to_observability",
        "rotation_only_accepted",
    }
    assert result.T_IMU_lidar is not None
    session0 = result.details["sessions"][0]
    assert "delta_v" in session0.get("pair_notes", []) or session0.get("pair_count", 0) >= 0
    # Phase-C fields appear when joint ran successfully on pairs.
    if session0.get("ok"):
        assert "gyro_bias_rad_s" in session0["joint"]


def test_signed_time_offset_refine_improves_or_keeps(tmp_path: Path):
    from imu_lidar.geometry import so3_exp
    from imu_lidar.time_offset import refine_time_offset_signed

    meta = generate_synthetic_session(tmp_path, delta_t_s=0.18, yaw_extrinsic_deg=20.0)
    imu = load_imu_samples(tmp_path / "imu.csv")
    frames = load_lidar_frames(tmp_path / "lidar")
    coarse = estimate_time_offset(imu, frames, search_s=0.5)
    r_true = so3_exp(np.deg2rad(np.array([2.0, -1.5, meta["yaw_extrinsic_deg"]])))
    refined = refine_time_offset_signed(
        imu,
        frames,
        delta_t_s=coarse.delta_t_s,
        R_IMU_lidar=r_true,
        search_s=0.08,
    )
    assert refined.ok
    # Must not drift farther from truth than the coarse estimate by a large margin.
    assert abs(refined.delta_t_s - meta["delta_t_s"]) <= abs(coarse.delta_t_s - meta["delta_t_s"]) + 0.01