calibration/imu_lidar/joint_optimizer.py

"""Joint extrinsic refinement: Phase-A rotation factors + Phase-C SE(3) IMU factors."""

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import dataclass, field
from typing import Any

import numpy as np
from scipy.optimize import least_squares

from .contracts import ImuSeries, MotionPair
from .geometry import make_transform, orthonormalize_rotation, so3_exp, so3_log
from .imu_preintegration import (
    apply_bias_jacobian_correction,
    apply_constant_bias_correction,
    preintegrate_gyro,
    preintegration_rotation_residual,
    residual_whiten_matrix,
)
from .observability import ObservabilityReport, analyze_observability
from .phase_a import phase_a_comparison_to_dict, solve_phase_a_comparison
from .rotation_handeye import select_strong_rotation_pairs

G_NORM = 9.80665


@dataclass(frozen=True)
class PhaseASessionResult:
    session_id: str
    pair_count: int
    gyro_bias0_rad_s: np.ndarray
    gyro_bias_rad_s: np.ndarray
    residual_rms_deg: float
    residual_median_deg: float
    residual_p95_deg: float
    outlier_fraction_gt_5deg: float
    accepted: bool
    included_in_final: bool


@dataclass(frozen=True)
class JointExtrinsicResult:
    T_IMU_lidar: np.ndarray
    translation_accepted: bool
    residual_rms_rot_deg: float
    residual_rms_trans_m: float
    observability: ObservabilityReport
    gyro_bias_rad_s: np.ndarray | None = None
    accel_bias_m_s2: np.ndarray | None = None
    gravity_m_s2: np.ndarray | None = None
    gyro_bias_rad_s_per_session: dict[str, np.ndarray] = field(default_factory=dict)
    phase_a_sessions: tuple[PhaseASessionResult, ...] = ()
    phase_a_accepted: bool = False
    phase_a_comparison: dict[str, Any] = field(default_factory=dict)
    notes: tuple[str, ...] = ()


def _pair_weight(pair: MotionPair) -> float:
    weight = float(pair.metadata.get("weight", 1.0))
    if not np.isfinite(weight) or weight <= 0:
        return 1.0
    return weight


def _pair_j_bg(pair: MotionPair) -> np.ndarray | None:
    raw = pair.metadata.get("J_bg")
    if raw is None:
        return None
    return np.asarray(raw, dtype=float).reshape(3, 3)


def _pair_cov(pair: MotionPair) -> np.ndarray:
    raw = pair.metadata.get("cov")
    if raw is None:
        sigma = float(pair.metadata.get("preint_sigma_rad", 1e-2))
        return np.eye(3) * max(sigma, 1e-4) ** 2
    return np.asarray(raw, dtype=float).reshape(3, 3)


def _corrected_delta_r(
    pair: MotionPair,
    delta_bias: np.ndarray,
    *,
    imu: ImuSeries | None,
    bias0: np.ndarray,
) -> np.ndarray:
    j_bg = _pair_j_bg(pair)
    if j_bg is not None:
        return apply_bias_jacobian_correction(pair.R_A, j_bg, delta_bias)
    if imu is not None and "t_i_imu_s" in pair.metadata and "t_j_imu_s" in pair.metadata:
        preint = preintegrate_gyro(
            imu.t_s,
            imu.gyro_rad_s,
            float(pair.metadata["t_i_imu_s"]),
            float(pair.metadata["t_j_imu_s"]),
            bias0 + delta_bias,
        )
        return preint.delta_R
    duration = float(pair.metadata.get("duration_s", max(pair.t_j_s - pair.t_i_s, 1e-3)))
    return apply_constant_bias_correction(pair.R_A, duration, delta_bias)


def _gravity_basis(g0: np.ndarray) -> np.ndarray:
    """Return 3×2 orthonormal basis spanning the plane orthogonal to ``g0``."""

    g = np.asarray(g0, dtype=float).reshape(3)
    n = np.linalg.norm(g)
    if n < 1e-9:
        g = np.array([0.0, 0.0, -G_NORM])
        n = G_NORM
    g = g / n
    axis = np.array([1.0, 0.0, 0.0]) if abs(g[0]) < 0.9 else np.array([0.0, 1.0, 0.0])
    e1 = np.cross(g, axis)
    e1 /= max(np.linalg.norm(e1), 1e-12)
    e2 = np.cross(g, e1)
    return np.column_stack([e1, e2])


def _gravity_from_params(xy: np.ndarray, g0: np.ndarray, basis: np.ndarray) -> np.ndarray:
    raw = np.asarray(g0, dtype=float).reshape(3) + basis @ np.asarray(xy, dtype=float).reshape(2)
    n = float(np.linalg.norm(raw))
    if n < 1e-9:
        return np.asarray(g0, dtype=float).reshape(3)
    return raw * (G_NORM / n)


def _lidar_to_imu_relative(r_x: np.ndarray, t_x: np.ndarray, r_b: np.ndarray, t_b: np.ndarray):
    """Map LiDAR relative pose to IMU: ``T_A = T_X T_B T_X^{-1}``."""

    r_a = orthonormalize_rotation(r_x @ r_b @ r_x.T)
    t_a = (np.eye(3) - r_a) @ t_x + r_x @ t_b
    return r_a, t_a


def _corrected_preint_quantities(
    pair: MotionPair,
    bg_i: np.ndarray,
    ba_i: np.ndarray,
    bg0: np.ndarray,
    ba0: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """First-order correct ΔR/Δv/Δp for keyframe biases vs preintegration biases."""

    dbg = np.asarray(bg_i, dtype=float).reshape(3) - np.asarray(bg0, dtype=float).reshape(3)
    dba = np.asarray(ba_i, dtype=float).reshape(3) - np.asarray(ba0, dtype=float).reshape(3)
    j_bg = pair.metadata.get("J_bg9")
    j_ba = pair.metadata.get("J_ba")
    delta_v0 = np.asarray(pair.metadata.get("delta_v", [0.0, 0.0, 0.0]), dtype=float).reshape(3)
    delta_p0 = (
        np.asarray(pair.t_A_m, dtype=float).reshape(3)
        if pair.t_A_m is not None
        else np.asarray(pair.metadata.get("delta_p", [0.0, 0.0, 0.0]), dtype=float).reshape(3)
    )
    if j_bg is None or j_ba is None:
        delta_r = apply_bias_jacobian_correction(
            pair.R_A,
            _pair_j_bg(pair) if _pair_j_bg(pair) is not None else np.zeros((3, 3)),
            dbg,
        )
        return delta_r, delta_v0, delta_p0
    j_bg_m = np.asarray(j_bg, dtype=float).reshape(9, 3)
    j_ba_m = np.asarray(j_ba, dtype=float).reshape(9, 3)
    delta_r = orthonormalize_rotation(pair.R_A @ so3_exp(j_bg_m[0:3] @ dbg))
    delta_v = delta_v0 + j_bg_m[3:6] @ dbg + j_ba_m[3:6] @ dba
    delta_p = delta_p0 + j_bg_m[6:9] @ dbg + j_ba_m[6:9] @ dba
    return delta_r, delta_v, delta_p


def _build_nav_rotations(
    keyframe_ids: list[int],
    id_to_idx: dict[int, int],
    consecutive_pairs: dict[tuple[int, int], MotionPair],
    r_x: np.ndarray,
    t_x: np.ndarray,
) -> list[np.ndarray]:
    """Chain IMU orientations; restart at session/gap boundaries (no cross-link)."""

    del id_to_idx
    rotations = [np.eye(3) for _ in keyframe_ids]
    for k in range(len(keyframe_ids) - 1):
        a = keyframe_ids[k]
        b = keyframe_ids[k + 1]
        pair = consecutive_pairs.get((a, b))
        if pair is None:
            # Missing link or new session: start a fresh nav chain.
            rotations[k + 1] = np.eye(3)
            continue
        t_b = np.zeros(3) if pair.t_B_m is None else np.asarray(pair.t_B_m, dtype=float)
        r_meas, _ = _lidar_to_imu_relative(r_x, t_x, pair.R_B, t_b)
        rotations[k + 1] = orthonormalize_rotation(rotations[k] @ r_meas)
    return rotations


def _solve_phase_c_se3(
    pairs: list[MotionPair],
    r_x: np.ndarray,
    *,
    gyro_bias_linearization: np.ndarray,
    gyro_bias_init: np.ndarray,
    gravity_init: np.ndarray,
    sigma_bg_rw: float = 1.0e-5,
    sigma_ba_rw: float = 1.0e-3,
    t_init: np.ndarray | None = None,
    t_prior: np.ndarray | None = None,
    t_prior_sigma_m: np.ndarray | float | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, float, float, list[str]]:
    """Keyframe IMU factor optimization for full SE(3)."""

    notes: list[str] = []
    usable = [pair for pair in pairs if pair.t_B_m is not None and "delta_v" in pair.metadata]
    if len(usable) < 3:
        notes.append("phase-C skipped: need pairs with full preintegration metadata")
        t0 = np.zeros(3) if t_init is None else np.asarray(t_init, dtype=float).reshape(3)
        return r_x, t0, gravity_init, gyro_bias_init, np.zeros(3), 1e9, 1e9, notes

    # Keyframes: group by session, sort each session by IMU time (no cross-session chain).
    stamp: dict[int, float] = {}
    kf_session: dict[int, str] = {}
    for pair in usable:
        stamp[pair.i] = float(pair.metadata.get("t_i_imu_s", pair.t_i_s))
        stamp[pair.j] = float(pair.metadata.get("t_j_imu_s", pair.t_j_s))
        kf_session[pair.i] = pair.session_id
        kf_session[pair.j] = pair.session_id
    session_ids = sorted(set(kf_session.values()))
    keyframe_ids: list[int] = []
    for sid in session_ids:
        local = [kid for kid, sess in kf_session.items() if sess == sid]
        local.sort(key=lambda kid: stamp[kid])
        keyframe_ids.extend(local)
    k_count = len(keyframe_ids)
    id_to_idx = {kid: idx for idx, kid in enumerate(keyframe_ids)}

    consecutive_pairs: dict[tuple[int, int], MotionPair] = {}
    for pair in usable:
        if kf_session.get(pair.i) != kf_session.get(pair.j):
            continue
        if id_to_idx[pair.j] == id_to_idx[pair.i] + 1:
            consecutive_pairs[(pair.i, pair.j)] = pair
    notes.append(
        f"phase-C multi-session graph: sessions={len(session_ids)}, "
        f"keyframes={k_count}, consecutive_links={len(consecutive_pairs)}"
    )

    g0 = np.asarray(gravity_init, dtype=float).reshape(3)
    if np.linalg.norm(g0) < 1e-6:
        g0 = np.array([0.0, 0.0, -G_NORM])
    g0 = g0 * (G_NORM / max(np.linalg.norm(g0), 1e-9))
    basis = _gravity_basis(g0)
    ba0 = np.zeros(3)
    bg0 = np.asarray(gyro_bias_linearization, dtype=float).reshape(3)
    bg_init = np.asarray(gyro_bias_init, dtype=float).reshape(3)

    # State: dθ(3), t(3), g_xy(2), v(3K), bg(3K), ba(3K)
    n_v = 3 * k_count
    n_b = 3 * k_count
    dim = 3 + 3 + 2 + n_v + n_b + n_b
    x0 = np.zeros(dim)
    t0 = np.zeros(3) if t_init is None else np.asarray(t_init, dtype=float).reshape(3)
    x0[3:6] = t0
    t_prior_vec = None if t_prior is None else np.asarray(t_prior, dtype=float).reshape(3)
    if t_prior_sigma_m is None:
        t_sigma = np.array([0.05, 0.05, 0.05], dtype=float)
    else:
        t_sigma = np.asarray(t_prior_sigma_m, dtype=float).reshape(-1)
        if t_sigma.size == 1:
            t_sigma = np.full(3, float(t_sigma[0]), dtype=float)
    # velocities start at 0; biases at prior
    for idx in range(k_count):
        x0[8 + n_v + 3 * idx : 8 + n_v + 3 * idx + 3] = bg_init

    whitened = []
    for pair in usable:
        cov9 = pair.metadata.get("cov9")
        if cov9 is None:
            cov = _pair_cov(pair)
            cov9_m = np.eye(9)
            cov9_m[0:3, 0:3] = cov
            cov9_m[3:6, 3:6] = np.eye(3) * 0.25
            cov9_m[6:9, 6:9] = np.eye(3) * 1.0
        else:
            cov9_m = np.asarray(cov9, dtype=float).reshape(9, 9)
        whitened.append(residual_whiten_matrix(cov9_m))

    def unpack(vec: np.ndarray):
        r_opt = orthonormalize_rotation(so3_exp(vec[0:3]) @ r_x)
        t_opt = vec[3:6]
        g_opt = _gravity_from_params(vec[6:8], g0, basis)
        base = 8
        vels = vec[base : base + n_v].reshape(k_count, 3)
        base += n_v
        bgs = vec[base : base + n_b].reshape(k_count, 3)
        base += n_b
        bas = vec[base : base + n_b].reshape(k_count, 3)
        return r_opt, t_opt, g_opt, vels, bgs, bas

    def residuals(vec: np.ndarray) -> np.ndarray:
        r_opt, t_opt, g_opt, vels, bgs, bas = unpack(vec)
        nav_r = _build_nav_rotations(keyframe_ids, id_to_idx, consecutive_pairs, r_opt, t_opt)
        out: list[np.ndarray] = []

        for pair, whiten in zip(usable, whitened):
            i_idx = id_to_idx[pair.i]
            j_idx = id_to_idx[pair.j]
            dt = float(pair.metadata.get("duration_s", pair.t_j_s - pair.t_i_s))
            dt = max(dt, 1e-3)
            delta_r, delta_v, delta_p = _corrected_preint_quantities(
                pair, bgs[i_idx], bas[i_idx], bg0, ba0
            )
            t_b = np.asarray(pair.t_B_m, dtype=float).reshape(3)
            r_meas, t_meas = _lidar_to_imu_relative(r_opt, t_opt, pair.R_B, t_b)
            r_i = nav_r[i_idx]
            v_i = vels[i_idx]
            v_j = vels[j_idx]

            err_r = so3_log(delta_r.T @ r_meas)
            err_v = v_j - v_i - g_opt * dt - r_i @ delta_v
            err_p = r_i @ (t_meas - delta_p) - v_i * dt - 0.5 * g_opt * (dt**2)
            err = np.concatenate([err_r, err_v, err_p])
            w = np.sqrt(_pair_weight(pair))
            out.append(w * (whiten @ err))

        # Bias random-walk between consecutive keyframes (same session only).
        for k in range(k_count - 1):
            a = keyframe_ids[k]
            b = keyframe_ids[k + 1]
            if kf_session.get(a) != kf_session.get(b):
                continue
            dt = max(stamp[b] - stamp[a], 1e-3)
            scale_g = 1.0 / (max(sigma_bg_rw, 1e-8) * np.sqrt(dt))
            scale_a = 1.0 / (max(sigma_ba_rw, 1e-8) * np.sqrt(dt))
            out.append(scale_g * (bgs[k + 1] - bgs[k]))
            out.append(scale_a * (bas[k + 1] - bas[k]))

        # Weak priors: first keyframe of each session + CAD/installation translation.
        for sid in session_ids:
            first = next(kid for kid in keyframe_ids if kf_session[kid] == sid)
            idx0 = id_to_idx[first]
            out.append(50.0 * (bgs[idx0] - bg_init))
            out.append(20.0 * bas[idx0])
        if t_prior_vec is not None:
            out.append((t_opt - t_prior_vec) / np.maximum(t_sigma, 1e-3))
        else:
            out.append(0.2 * t_opt)  # soft |t|~0 prior when no CAD prior
        return np.concatenate(out)

    # Cap evaluations: Phase-C is high-dimensional; synthetic ICP already dominates runtime.
    opt = least_squares(residuals, x0, loss="huber", f_scale=0.05, max_nfev=80)
    r_opt, t_opt, g_opt, vels, bgs, bas = unpack(opt.x)

    rot_errs = []
    trans_errs = []
    nav_r = _build_nav_rotations(keyframe_ids, id_to_idx, consecutive_pairs, r_opt, t_opt)
    for pair in usable:
        i_idx = id_to_idx[pair.i]
        j_idx = id_to_idx[pair.j]
        dt = max(float(pair.metadata.get("duration_s", pair.t_j_s - pair.t_i_s)), 1e-3)
        delta_r, delta_v, delta_p = _corrected_preint_quantities(
            pair, bgs[i_idx], bas[i_idx], bg0, ba0
        )
        t_b = np.asarray(pair.t_B_m, dtype=float).reshape(3)
        r_meas, t_meas = _lidar_to_imu_relative(r_opt, t_opt, pair.R_B, t_b)
        r_i = nav_r[i_idx]
        err_r = so3_log(delta_r.T @ r_meas)
        err_p = r_i @ (t_meas - delta_p) - vels[i_idx] * dt - 0.5 * g_opt * (dt**2)
        rot_errs.append(np.degrees(np.linalg.norm(err_r)))
        trans_errs.append(float(np.linalg.norm(err_p)))
        del delta_v, j_idx

    rot_rms = float(np.sqrt(np.mean(np.square(rot_errs)))) if rot_errs else 1e9
    trans_rms = float(np.sqrt(np.mean(np.square(trans_errs)))) if trans_errs else 1e9
    bg_mean = np.mean(bgs, axis=0)
    ba_mean = np.mean(bas, axis=0)
    notes.append(
        "phase-C SE3 (Δv/Δp + g + keyframe v/bias RW): "
        f"keyframes={k_count}, pairs={len(usable)}, "
        f"|t|={float(np.linalg.norm(t_opt)):.3f} m, "
        f"|g|={float(np.linalg.norm(g_opt)):.3f}, "
        f"trans_rms={trans_rms:.3f} m"
    )
    return r_opt, t_opt, g_opt, bg_mean, ba_mean, rot_rms, trans_rms, notes


def _pair_gyro_bias0(pair: MotionPair, fallback: np.ndarray) -> np.ndarray:
    raw = pair.metadata.get("gyro_bias0_rad_s")
    if raw is None:
        return np.asarray(fallback, dtype=float).reshape(3)
    return np.asarray(raw, dtype=float).reshape(3)


def _phase_a_bias_bases(
    pairs: list[MotionPair],
    *,
    gyro_bias_rad_s: np.ndarray | None,
    gyro_bias_rad_s_by_session: Mapping[str, np.ndarray] | None,
) -> dict[str, np.ndarray]:
    session_ids = sorted({pair.session_id for pair in pairs})
    scalar = None
    if gyro_bias_rad_s is not None:
        scalar = np.asarray(gyro_bias_rad_s, dtype=float).reshape(3)
    supplied = {} if gyro_bias_rad_s_by_session is None else gyro_bias_rad_s_by_session
    bases: dict[str, np.ndarray] = {}
    for sid in session_ids:
        if sid in supplied:
            bases[sid] = np.asarray(supplied[sid], dtype=float).reshape(3)
            continue
        pair = next(
            (
                item
                for item in pairs
                if item.session_id == sid and "gyro_bias0_rad_s" in item.metadata
            ),
            None,
        )
        if pair is not None:
            bases[sid] = np.asarray(pair.metadata["gyro_bias0_rad_s"], dtype=float).reshape(3)
        elif scalar is not None:
            bases[sid] = scalar.copy()
        else:
            bases[sid] = np.zeros(3)
    return bases


def _rotation_distribution(errs_deg: list[float]) -> tuple[float, float, float, float, bool]:
    if not errs_deg:
        return 1e9, 1e9, 1e9, 1.0, False
    errs = np.asarray(errs_deg, dtype=float)
    rms = float(np.sqrt(np.mean(errs**2)))
    median = float(np.median(errs))
    p95 = float(np.percentile(errs, 95.0))
    outlier_fraction = float(np.mean(errs > 5.0))
    accepted = (
        len(errs) >= 3
        and rms < 1.5
        and median < 0.5
        and p95 < 1.5
        and outlier_fraction <= 0.005
    )
    return rms, median, p95, outlier_fraction, accepted


def _solve_phase_a_rotation(
    pairs: list[MotionPair],
    r_seed: np.ndarray,
    *,
    bias_bases: Mapping[str, np.ndarray],
    imu: ImuSeries | None,
    bias_prior_sigma_rad_s: float,
    preexcluded_session_ids: set[str] | None = None,
) -> tuple[
    np.ndarray,
    dict[str, np.ndarray],
    tuple[PhaseASessionResult, ...],
    list[MotionPair],
    float,
    bool,
    list[str],
]:
    notes: list[str] = []
    all_session_ids = sorted({pair.session_id for pair in pairs})
    prior_w = 1.0 / max(bias_prior_sigma_rad_s, 1e-4)

    def optimize(
        active_pairs: list[MotionPair],
        r0: np.ndarray,
        bias_seed: Mapping[str, np.ndarray],
    ) -> tuple[np.ndarray, dict[str, np.ndarray]]:
        session_ids = sorted({pair.session_id for pair in active_pairs})
        session_index = {sid: index for index, sid in enumerate(session_ids)}
        whiten = [residual_whiten_matrix(_pair_cov(pair)) for pair in active_pairs]
        x0 = np.zeros(3 + 3 * len(session_ids))
        for sid, index in session_index.items():
            x0[3 + 3 * index : 6 + 3 * index] = np.asarray(bias_seed[sid], dtype=float)

        def residual(vec: np.ndarray) -> np.ndarray:
            r_opt = orthonormalize_rotation(so3_exp(vec[:3]) @ r0)
            out: list[np.ndarray] = []
            for pair, sqrt_info in zip(active_pairs, whiten):
                index = session_index[pair.session_id]
                bias = vec[3 + 3 * index : 6 + 3 * index]
                base = _pair_gyro_bias0(pair, bias_bases[pair.session_id])
                delta_r = _corrected_delta_r(
                    pair, bias - base, imu=imu, bias0=base
                )
                out.append(
                    sqrt_info
                    @ preintegration_rotation_residual(delta_r, r_opt, pair.R_B)
                )
            for sid, index in session_index.items():
                bias = vec[3 + 3 * index : 6 + 3 * index]
                out.append(prior_w * (bias - bias_bases[sid]))
            return np.concatenate(out)

        opt = least_squares(residual, x0, loss="huber", f_scale=1.0, max_nfev=200)
        r_opt = orthonormalize_rotation(so3_exp(opt.x[:3]) @ r0)
        biases = {
            sid: opt.x[3 + 3 * index : 6 + 3 * index].copy()
            for sid, index in session_index.items()
        }
        return r_opt, biases

    def summarize(
        r_opt: np.ndarray,
        biases: Mapping[str, np.ndarray],
        included: set[str],
    ) -> tuple[PhaseASessionResult, ...]:
        results: list[PhaseASessionResult] = []
        for sid in all_session_ids:
            local_pairs = [pair for pair in pairs if pair.session_id == sid]
            bias = np.asarray(biases.get(sid, bias_bases[sid]), dtype=float).reshape(3)
            errs: list[float] = []
            for pair in local_pairs:
                base = _pair_gyro_bias0(pair, bias_bases[sid])
                delta_r = _corrected_delta_r(pair, bias - base, imu=imu, bias0=base)
                err = preintegration_rotation_residual(delta_r, r_opt, pair.R_B)
                errs.append(float(np.degrees(np.linalg.norm(err))))
            rms, median, p95, outlier, accepted = _rotation_distribution(errs)
            results.append(
                PhaseASessionResult(
                    session_id=sid,
                    pair_count=len(local_pairs),
                    gyro_bias0_rad_s=np.asarray(bias_bases[sid], dtype=float),
                    gyro_bias_rad_s=bias,
                    residual_rms_deg=rms,
                    residual_median_deg=median,
                    residual_p95_deg=p95,
                    outlier_fraction_gt_5deg=outlier,
                    accepted=accepted,
                    included_in_final=sid in included,
                )
            )
        return tuple(results)

    if not pairs:
        return r_seed, dict(bias_bases), (), [], 1e9, False, ["no pairs for phase-A"]

    r_first, biases_first = optimize(pairs, r_seed, bias_bases)
    first = summarize(r_first, biases_first, set(all_session_ids))
    accepted_ids = {item.session_id for item in first if item.accepted}
    preexcluded = set() if preexcluded_session_ids is None else set(preexcluded_session_ids)
    accepted_ids -= preexcluded
    active_ids = set(all_session_ids)
    r_final = r_first
    biases_final = dict(biases_first)
    if preexcluded and not accepted_ids:
        active_ids = set()
        notes.append(f"phase-A pre-gate excluded all sessions: {sorted(preexcluded)}")
    elif accepted_ids and accepted_ids != active_ids:
        active_ids = accepted_ids
        active_pairs = [pair for pair in pairs if pair.session_id in active_ids]
        r_final, active_biases = optimize(active_pairs, r_first, biases_first)
        biases_final.update(active_biases)
        excluded = sorted(set(all_session_ids) - active_ids)
        notes.append(f"phase-A excluded sessions after local/pre residual gate: {excluded}")
    active_pairs = [pair for pair in pairs if pair.session_id in active_ids]
    final = summarize(r_final, biases_final, active_ids)
    active_results = [item for item in final if item.included_in_final]
    global_errs: list[float] = []
    for pair in active_pairs:
        bias = biases_final[pair.session_id]
        base = _pair_gyro_bias0(pair, bias_bases[pair.session_id])
        delta_r = _corrected_delta_r(pair, bias - base, imu=imu, bias0=base)
        err = preintegration_rotation_residual(delta_r, r_final, pair.R_B)
        global_errs.append(float(np.degrees(np.linalg.norm(err))))
    rot_rms, _, _, _, global_ok = _rotation_distribution(global_errs)
    accepted = bool(active_results and global_ok and all(item.accepted for item in active_results))
    notes.append(
        f"phase-A session-local bias refine: sessions={len(active_ids)}/{len(all_session_ids)}, "
        f"pairs={len(active_pairs)}, rms={rot_rms:.3f} deg"
    )
    return r_final, biases_final, final, active_pairs, rot_rms, accepted, notes


def _solve_joint_extrinsic_legacy(
    pairs: list[MotionPair] | tuple[MotionPair, ...],
    r_x: np.ndarray,
    *,
    force_rotation_only: bool = False,
    imu: ImuSeries | None = None,
    delta_t_s: float = 0.0,
    gyro_bias_rad_s: np.ndarray | None = None,
    gravity_init_m_s2: np.ndarray | None = None,
    gyro_bias_rad_s_by_session: Mapping[str, np.ndarray] | None = None,
    time_offset_s_by_session: Mapping[str, float] | None = None,
    bias_prior_sigma_rad_s: float = 0.02,
    enable_phase_c: bool | None = None,
    t_init_m: np.ndarray | None = None,
    t_prior_m: np.ndarray | None = None,
    t_prior_sigma_m: np.ndarray | float | None = None,
) -> JointExtrinsicResult:
    """Refine extrinsic using Phase-A whitened rotation factors, optional Phase-C SE(3)."""

    del delta_t_s  # reserved for future SE(3) time coupling
    if enable_phase_c is None:
        enable_phase_c = not force_rotation_only

    usable = [pair for pair in pairs if pair.t_B_m is not None]
    observability = analyze_observability(usable, r_x)
    notes = list(observability.notes)

    r = orthonormalize_rotation(np.asarray(r_x, dtype=float))
    bias0 = np.zeros(3) if gyro_bias_rad_s is None else np.asarray(gyro_bias_rad_s, dtype=float).reshape(3)
    t_seed = None if t_init_m is None else np.asarray(t_init_m, dtype=float).reshape(3)
    weights = np.asarray([_pair_weight(pair) for pair in usable], dtype=float)
    whitens = [residual_whiten_matrix(_pair_cov(pair)) for pair in usable]
    prior_w = 1.0 / max(bias_prior_sigma_rad_s, 1e-4)

    def rotation_residuals(r_opt: np.ndarray, delta_bias: np.ndarray) -> np.ndarray:
        residuals = []
        for pair, whiten in zip(usable, whitens):
            delta_r = _corrected_delta_r(pair, delta_bias, imu=imu, bias0=bias0)
            err = preintegration_rotation_residual(delta_r, r_opt, pair.R_B)
            residuals.append(whiten @ err)
        residuals.append(prior_w * delta_bias)
        return np.concatenate(residuals) if residuals else np.zeros(0)

    def residual_rot_bias(vec: np.ndarray) -> np.ndarray:
        r_opt = orthonormalize_rotation(so3_exp(vec[:3]) @ r)
        return rotation_residuals(r_opt, vec[3:])

    if usable:
        opt = least_squares(
            residual_rot_bias,
            np.zeros(6),
            loss="huber",
            f_scale=1.0,
            max_nfev=200,
        )
        r = orthonormalize_rotation(so3_exp(opt.x[:3]) @ r)
        delta_bias = opt.x[3:]
        bias_out = bias0 + delta_bias
        notes.append(
            "phase-A joint refine (single Σ whitening + J_bg): "
            f"|δb|={float(np.linalg.norm(delta_bias)):.3e} rad/s, "
            f"pairs={len(usable)}"
        )
    else:
        bias_out = bias0
        delta_bias = np.zeros(3)
        notes.append("no pairs for joint refine")

    rot_errs = []
    for pair in usable:
        delta_r = _corrected_delta_r(pair, delta_bias, imu=imu, bias0=bias0)
        err = preintegration_rotation_residual(delta_r, r, pair.R_B)
        rot_errs.append(np.degrees(np.linalg.norm(err)))
    rot_rms = float(np.sqrt(np.mean(np.square(rot_errs)))) if rot_errs else 1e9

    t = np.zeros(3) if t_seed is None else t_seed.copy()
    translation_accepted = False
    trans_rms = 1e9
    gravity_out: np.ndarray | None = None
    accel_bias_out: np.ndarray | None = None

    if gravity_init_m_s2 is None:
        gravity_init = np.array([0.0, 0.0, -G_NORM])
    else:
        gravity_init = np.asarray(gravity_init_m_s2, dtype=float).reshape(3)

    if t_prior_m is not None:
        notes.append(
            "using CAD/installation translation prior "
            f"t={np.asarray(t_prior_m, dtype=float).reshape(3).tolist()}"
        )

    if (
        enable_phase_c
        and not force_rotation_only
        and observability.translation_observable
        and observability.rotation_observable
        and len(usable) >= 5
    ):
        r, t, gravity_out, bias_out, accel_bias_out, rot_rms, trans_rms, c_notes = _solve_phase_c_se3(
            usable,
            r,
            gyro_bias_linearization=bias0,
            gyro_bias_init=bias_out,
            gravity_init=gravity_init,
            t_init=t_seed if t_seed is not None else t_prior_m,
            t_prior=t_prior_m,
            t_prior_sigma_m=t_prior_sigma_m,
        )
        notes.extend(c_notes)
        translation_accepted = bool(trans_rms < 0.75 and np.linalg.norm(t) > 1e-4)
        if not translation_accepted:
            # Prefer CAD prior over silent zero when motion SE3 is rejected.
            if t_prior_m is not None:
                t = np.asarray(t_prior_m, dtype=float).reshape(3)
                notes.append(
                    "phase-C translation residual/gate failed; CAD translation is reported "
                    "as a prior only and is not accepted as calibration"
                )
            else:
                notes.append("phase-C translation residual/gate failed; keeping translation at zero")
                t = np.zeros(3)
    elif (
        not force_rotation_only
        and observability.translation_observable
        and observability.rotation_observable
        and len(usable) >= 5
    ):
        # Legacy hand-eye translation fallback when Phase-C metadata missing.
        def residual_se3(vec: np.ndarray) -> np.ndarray:
            r_opt = orthonormalize_rotation(so3_exp(vec[:3]) @ r)
            t_opt = vec[3:]
            residuals = []
            for pair, weight, whiten in zip(usable, weights, whitens):
                delta_r = _corrected_delta_r(pair, delta_bias, imu=imu, bias0=bias0)
                residuals.append(
                    np.sqrt(weight) * (whiten @ preintegration_rotation_residual(delta_r, r_opt, pair.R_B))
                )
                pred = (pair.R_A - np.eye(3)) @ t_opt
                meas = r_opt @ np.asarray(pair.t_B_m, dtype=float)
                residuals.append(np.sqrt(weight) * (pred - meas))
            if t_prior_m is not None:
                sigma = np.asarray(t_prior_sigma_m if t_prior_sigma_m is not None else 0.05, dtype=float)
                if sigma.size == 1:
                    sigma = np.full(3, float(sigma), dtype=float)
                residuals.append((t_opt - np.asarray(t_prior_m, dtype=float).reshape(3)) / np.maximum(sigma, 1e-3))
            return np.concatenate(residuals)

        x_se3 = np.zeros(6)
        if t_seed is not None:
            x_se3[3:] = t_seed
        elif t_prior_m is not None:
            x_se3[3:] = np.asarray(t_prior_m, dtype=float).reshape(3)
        opt_t = least_squares(residual_se3, x_se3, loss="huber", f_scale=0.05, max_nfev=200)
        r = orthonormalize_rotation(so3_exp(opt_t.x[:3]) @ r)
        t = opt_t.x[3:]
        rot_errs = []
        trans_errs = []
        for pair in usable:
            delta_r = _corrected_delta_r(pair, delta_bias, imu=imu, bias0=bias0)
            rot_errs.append(np.degrees(np.linalg.norm(preintegration_rotation_residual(delta_r, r, pair.R_B))))
            pred = (pair.R_A - np.eye(3)) @ t
            meas = r @ np.asarray(pair.t_B_m, dtype=float)
            trans_errs.append(np.linalg.norm(pred - meas))
        rot_rms = float(np.sqrt(np.mean(np.square(rot_errs))))
        trans_rms = float(np.sqrt(np.mean(np.square(trans_errs))))
        translation_accepted = trans_rms < 0.5
        notes.append(f"legacy translation refine rms={trans_rms:.3f} m")
        if not translation_accepted:
            notes.append("translation residual too large; keeping translation at zero")
            t = np.zeros(3)
    elif not force_rotation_only and t_prior_m is not None:
        t = np.asarray(t_prior_m, dtype=float).reshape(3)
        translation_accepted = False
        notes.append(
            "SE3 motion solve gated off; CAD translation is reported as a prior only "
            "and is not accepted as calibration"
        )
    else:
        notes.append("rotation-only extrinsic returned (phase-A; phase-C SE3 gated off)")

    return JointExtrinsicResult(
        T_IMU_lidar=make_transform(t, r),
        translation_accepted=bool(translation_accepted and np.linalg.norm(t) > 0),
        residual_rms_rot_deg=rot_rms,
        residual_rms_trans_m=trans_rms,
        observability=observability,
        gyro_bias_rad_s=np.asarray(bias_out, dtype=float),
        accel_bias_m_s2=None if accel_bias_out is None else np.asarray(accel_bias_out, dtype=float),
        gravity_m_s2=None if gravity_out is None else np.asarray(gravity_out, dtype=float),
        notes=tuple(notes),
    )


def solve_joint_extrinsic(
    pairs: list[MotionPair] | tuple[MotionPair, ...],
    r_x: np.ndarray,
    *,
    force_rotation_only: bool = False,
    imu: ImuSeries | None = None,
    delta_t_s: float = 0.0,
    gyro_bias_rad_s: np.ndarray | None = None,
    gyro_bias_rad_s_by_session: Mapping[str, np.ndarray] | None = None,
    time_offset_s_by_session: Mapping[str, float] | None = None,
    preexcluded_session_ids: set[str] | None = None,
    gravity_init_m_s2: np.ndarray | None = None,
    bias_prior_sigma_rad_s: float = 0.002,
    rotation_prior: np.ndarray | None = None,
    rotation_prior_sigma_deg: float = 15.0,
    phase_a_yaw_std_max_deg: float = 0.5,
    phase_a_loo_yaw_range_max_deg: float = 1.0,
    phase_a_data_prior_difference_max_deg: float = 1.0,
    run_phase_a_leave_one_out: bool = True,
    phase_a_progress_callback: (
        Callable[[str, dict[str, Any]], None] | None
    ) = None,
    enable_phase_c: bool | None = None,
    t_init_m: np.ndarray | None = None,
    t_prior_m: np.ndarray | None = None,
    t_prior_sigma_m: np.ndarray | float | None = None,
) -> JointExtrinsicResult:
    """Run the corrected session-aware Phase-A and gate unfinished SE(3) stages."""

    del gravity_init_m_s2, t_init_m, t_prior_sigma_m, imu, r_x
    usable_input = [pair for pair in pairs if pair.t_B_m is not None]
    bias_bases = _phase_a_bias_bases(
        usable_input,
        gyro_bias_rad_s=gyro_bias_rad_s,
        gyro_bias_rad_s_by_session=gyro_bias_rad_s_by_session,
    )
    comparison = solve_phase_a_comparison(
        usable_input,
        gyro_bias_rad_s_by_session=bias_bases,
        rotation_prior=rotation_prior,
        rotation_prior_sigma_deg=rotation_prior_sigma_deg,
        preexcluded_session_ids=preexcluded_session_ids,
        bias_prior_sigma_rad_s=bias_prior_sigma_rad_s,
        yaw_std_max_deg=phase_a_yaw_std_max_deg,
        leave_one_out_yaw_range_max_deg=(
            phase_a_loo_yaw_range_max_deg
        ),
        data_prior_difference_max_deg=(
            phase_a_data_prior_difference_max_deg
        ),
        run_leave_one_out=run_phase_a_leave_one_out,
        progress_callback=phase_a_progress_callback,
    )
    primary = comparison.session_bg_data_only
    r = primary.R_IMU_lidar
    biases = primary.gyro_bias_rad_s_per_session
    rot_rms = primary.residual_rms_deg
    phase_a_accepted = comparison.accepted
    notes = list(comparison.notes)
    notes.append(
        "phase-A primary=A1_session_bg_data_only; "
        f"A0 RPY={comparison.fixed_bg_data_only.rpy_deg_xyz.tolist()}, "
        f"A1 RPY={primary.rpy_deg_xyz.tolist()}, "
        "A2 RPY="
        f"{comparison.session_bg_with_rotation_prior.rpy_deg_xyz.tolist()}"
    )
    notes.append(
        f"phase-A marginal yaw_std={comparison.marginal_observability.yaw_std_deg:.3f} deg, "
        f"LOO yaw range={comparison.leave_one_out_yaw_range_deg:.3f} deg"
    )

    session_results_list: list[PhaseASessionResult] = [
        PhaseASessionResult(
            session_id=item.session_id,
            pair_count=item.pair_count,
            gyro_bias0_rad_s=item.gyro_bias0_rad_s,
            gyro_bias_rad_s=item.gyro_bias_rad_s,
            residual_rms_deg=item.residual_rms_deg,
            residual_median_deg=item.residual_median_deg,
            residual_p95_deg=item.residual_p95_deg,
            outlier_fraction_gt_5deg=item.outlier_fraction_gt_5deg,
            accepted=item.accepted,
            included_in_final=True,
        )
        for item in primary.sessions
    ]
    preexcluded = (
        set()
        if preexcluded_session_ids is None
        else set(preexcluded_session_ids)
    )
    strong_all = select_strong_rotation_pairs(usable_input)
    for session_id in sorted(preexcluded):
        local_pairs = [
            pair for pair in strong_all if pair.session_id == session_id
        ]
        errors = [
            float(
                np.degrees(
                    np.linalg.norm(
                        preintegration_rotation_residual(
                            pair.R_A, r, pair.R_B
                        )
                    )
                )
            )
            for pair in local_pairs
        ]
        rms, median, p95, outlier, accepted = _rotation_distribution(
            errors
        )
        base = np.asarray(
            bias_bases.get(session_id, np.zeros(3)), dtype=float
        ).reshape(3)
        session_results_list.append(
            PhaseASessionResult(
                session_id=session_id,
                pair_count=len(local_pairs),
                gyro_bias0_rad_s=base,
                gyro_bias_rad_s=base,
                residual_rms_deg=rms,
                residual_median_deg=median,
                residual_p95_deg=p95,
                outlier_fraction_gt_5deg=outlier,
                accepted=accepted,
                included_in_final=False,
            )
        )
    session_results = tuple(
        sorted(session_results_list, key=lambda item: item.session_id)
    )
    usable = [
        pair
        for pair in strong_all
        if pair.session_id not in preexcluded
    ]
    base_observability = analyze_observability(usable, r)
    marginal = comparison.marginal_observability
    observability = ObservabilityReport(
        rotation_observable=bool(
            marginal.rank == 3
            and marginal.yaw_std_deg <= phase_a_yaw_std_max_deg
        ),
        translation_observable=base_observability.translation_observable,
        condition_rotation=marginal.condition,
        condition_translation=base_observability.condition_translation,
        notes=tuple(
            list(marginal.notes)
            + list(base_observability.notes)
        ),
    )
    notes.extend(observability.notes)
    if time_offset_s_by_session is None:
        notes.append(
            f"legacy scalar time offset fixed during pair construction: {float(delta_t_s):.6f}s"
        )
    else:
        fixed_offsets = {
            str(sid): float(value) for sid, value in time_offset_s_by_session.items()
        }
        notes.append(
            f"time offsets fixed during pair construction (not optimized): {fixed_offsets}"
        )

    for item in session_results:
        notes.append(
            f"phase-A session {item.session_id}: included={item.included_in_final}, "
            f"pairs={item.pair_count}, rms={item.residual_rms_deg:.3f} deg, "
            f"p95={item.residual_p95_deg:.3f} deg, "
            f"|bias-bias0|={float(np.linalg.norm(item.gyro_bias_rad_s - item.gyro_bias0_rad_s)):.3e}"
        )

    phase_c_requested = (not force_rotation_only) if enable_phase_c is None else bool(enable_phase_c)
    t = np.zeros(3)
    if not force_rotation_only:
        if phase_c_requested:
            notes.append(
                "phase-B/C gated off: session-aware translation/gravity/navigation "
                "states are not implemented yet"
            )
        else:
            notes.append("phase-C disabled; translation is not accepted")
        if t_prior_m is not None:
            t = np.asarray(t_prior_m, dtype=float).reshape(3)
            notes.append(
                "CAD translation is reported as a prior only and is not accepted as calibration"
            )
    else:
        notes.append("rotation-only extrinsic returned after corrected phase-A")

    single_bias = None
    if len(biases) == 1:
        single_bias = np.asarray(next(iter(biases.values())), dtype=float)
    return JointExtrinsicResult(
        T_IMU_lidar=make_transform(t, r),
        translation_accepted=False,
        residual_rms_rot_deg=rot_rms,
        residual_rms_trans_m=1e9,
        observability=observability,
        gyro_bias_rad_s=single_bias,
        accel_bias_m_s2=None,
        gravity_m_s2=None,
        gyro_bias_rad_s_per_session={
            sid: np.asarray(value, dtype=float) for sid, value in biases.items()
        },
        phase_a_sessions=session_results,
        phase_a_accepted=phase_a_accepted,
        phase_a_comparison=phase_a_comparison_to_dict(comparison),
        notes=tuple(notes),
    )