calibration/imu_lidar/joint_optimizer.py

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

from __future__ import annotations

from dataclasses import dataclass

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

G_NORM = 9.80665


@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
    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 in the first-keyframe nav frame using LiDAR+extrinsic."""

    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:
            rotations[k + 1] = rotations[k]
            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)
    # Ensure list indexed by id_to_idx
    del id_to_idx
    return rotations


def _solve_phase_c_se3(
    pairs: list[MotionPair],
    r_x: np.ndarray,
    *,
    gyro_bias0: np.ndarray,
    gravity_init: np.ndarray,
    sigma_bg_rw: float = 1.0e-5,
    sigma_ba_rw: float = 1.0e-3,
) -> 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")
        return r_x, np.zeros(3), gravity_init, gyro_bias0, np.zeros(3), 1e9, 1e9, notes

    # Unique keyframes sorted by IMU time.
    stamp: dict[int, float] = {}
    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))
    keyframe_ids = sorted(stamp.keys(), key=lambda kid: stamp[kid])
    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 id_to_idx[pair.j] == id_to_idx[pair.i] + 1:
            consecutive_pairs[(pair.i, pair.j)] = pair

    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_bias0, 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)
    # 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] = bg0

    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.
        for k in range(k_count - 1):
            dt = max(stamp[keyframe_ids[k + 1]] - stamp[keyframe_ids[k]], 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 biases and translation magnitude.
        out.append(50.0 * (bgs[0] - bg0))
        out.append(20.0 * bas[0])
        out.append(0.2 * t_opt)  # soft |t| prior ~ meters
        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 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,
    gravity_init_m_s2: np.ndarray | None = None,
    bias_prior_sigma_rad_s: float = 0.02,
    enable_phase_c: bool | 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)
    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, weight, whiten in zip(usable, weights, 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(np.sqrt(weight) * (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=np.deg2rad(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 (Σ-whitened + J_bg): "
            f"|δb|={float(np.linalg.norm(delta_bias)):.3e} rad/s, "
            f"weighted 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)
    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 (
        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_bias0=bias_out,
            gravity_init=gravity_init,
        )
        notes.extend(c_notes)
        translation_accepted = bool(trans_rms < 0.75 and np.linalg.norm(t) > 1e-4)
        if not translation_accepted:
            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))
            return np.concatenate(residuals)

        opt_t = least_squares(residual_se3, np.zeros(6), 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)
    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=0.0 if not translation_accepted else 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),
    )