"""Optional exact-order CPU kernel for continuous-time pose/world coordinates."""

from __future__ import annotations

import math
import sys
from collections import OrderedDict
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from contextlib import suppress
from dataclasses import dataclass, field
from importlib.metadata import version
from itertools import pairwise
from threading import Lock
from typing import TYPE_CHECKING, TypeAlias, cast

import numpy as np
from numba import njit
from numpy.typing import NDArray

from calibrex.core.geometry import SE3, Vector3
from calibrex.data.odometry_track import OdometryTrack, _PreparedInterval

if TYPE_CHECKING:
    from calibrex.evaluation.packed_lidar_observations import _PlaneColumns
    from calibrex.graph.lidar_point_to_plane import _ObservationArrays

FloatArray: TypeAlias = NDArray[np.float64]
IntArray: TypeAlias = NDArray[np.int64]
Quaternion: TypeAlias = tuple[float, float, float, float]


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _association_choices(
    world: FloatArray, left: IntArray, right: IntArray,
    lc: FloatArray, lr: FloatArray, lv: NDArray[np.bool_],
    rc: FloatArray, rr: FloatArray, rv: NDArray[np.bool_], hybrid: bool,
) -> tuple[NDArray[np.bool_], NDArray[np.bool_], IntArray, bool]:
    n = len(world)
    prefer_right = np.zeros(n, dtype=np.bool_)
    valid = np.zeros(n, dtype=np.bool_)
    ties = np.empty(n, dtype=np.int64)
    tie_count = 0
    finite = True
    for row in range(n):
        lp, rp = max(left[row], 0), max(right[row], 0)
        present_left = left[row] >= 0 and lv[lp]
        if not hybrid:
            valid[row] = present_left
            continue
        present_right = right[row] >= 0 and rv[rp]
        a = abs(lr[lp, 0]*(world[row, 0]-lc[lp, 0])
                + lr[lp, 1]*(world[row, 1]-lc[lp, 1])
                + lr[lp, 2]*(world[row, 2]-lc[lp, 2]))
        b = abs(rr[rp, 0]*(world[row, 0]-rc[rp, 0])
                + rr[rp, 1]*(world[row, 1]-rc[rp, 1])
                + rr[rp, 2]*(world[row, 2]-rc[rp, 2]))
        finite = finite and math.isfinite(a) and math.isfinite(b)
        valid[row] = present_left or present_right
        prefer_right[row] = present_right and (not present_left or b < a)
        # Equal centroids also have equal math.dist, including signed zero.
        # Preserve the first (uniform) plane; repair other exact ties in Python.
        if (present_left and present_right and a == b
            and not (lc[lp, 0] == rc[rp, 0] and lc[lp, 1] == rc[rp, 1]
                     and lc[lp, 2] == rc[rp, 2])):
            ties[tie_count] = row
            tie_count += 1
    return prefer_right, valid, ties[:tie_count], finite


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _gather_observation_rows(
    points: FloatArray, indices: IntArray, poses: FloatArray,
    left: IntArray, right: IntArray, lc: FloatArray, ln: FloatArray,
    rc: FloatArray, rn: FloatArray, prefer_right: NDArray[np.bool_],
    valid: NDArray[np.bool_],
) -> tuple[FloatArray, FloatArray, FloatArray, FloatArray, FloatArray]:
    count = np.count_nonzero(valid)
    xyz = np.empty((count, 3), dtype=np.float64)
    centers = np.empty((count, 3), dtype=np.float64)
    normals = np.empty((count, 3), dtype=np.float64)
    rows = np.empty((count, 7), dtype=np.float64)
    weights = np.ones(count, dtype=np.float64)
    target = 0
    for row in range(len(valid)):
        if not valid[row]:
            continue
        position = right[row] if prefer_right[row] else left[row]
        for axis in range(3):
            xyz[target, axis] = points[indices[row], axis]
            centers[target, axis] = rc[position, axis] if prefer_right[row] else lc[position, axis]
            normals[target, axis] = rn[position, axis] if prefer_right[row] else ln[position, axis]
        for axis in range(7):
            rows[target, axis] = poses[row, axis]
        target += 1
    return xyz, centers, normals, rows, weights


def _pack_rig_observations(
    points: FloatArray, indices: IntArray, poses: FloatArray, world: FloatArray,
    matches: Sequence[IntArray], columns: Sequence[_PlaneColumns],
) -> _ObservationArrays | None:
    """Fuse ordered native scores and gathers; retain Python math.dist for exact ties."""
    from calibrex.graph.lidar_point_to_plane import _ObservationArrays

    if (len(matches) not in (1, 2) or len(matches) != len(columns)
        or world.shape != (len(world), 3) or poses.shape != (len(world), 7)
        or indices.shape != (len(world),)
        or np.any(indices >= len(points)) or np.any(indices < -len(points))
        or any(position.shape != (len(world),) or len(column.centroids) == 0
               or np.any(position >= len(column.centroids))
               for position, column in zip(matches, columns, strict=True))):
        return None
    left, right = matches[0], matches[-1]
    lc, rc = columns[0], columns[-1]
    prefer_right, valid, ties, finite = _association_choices(
        world, left, right, lc.centroids, lc.raw_normals, lc.valid,
        rc.centroids, rc.raw_normals, rc.valid, len(matches) == 2,
    )
    if not finite:
        return None
    for row in ties:
        point = tuple(world[row].tolist())
        prefer_right[row] = (math.dist(point, rc.centroids[right[row]].tolist())
                             < math.dist(point, lc.centroids[left[row]].tolist()))
    arrays = _gather_observation_rows(
        points, indices, poses, left, right, lc.centroids, lc.normals,
        rc.centroids, rc.normals, prefer_right, valid,
    )
    for array in arrays:
        array.flags.writeable = False
    xyz, centers, normals, rows, weights = arrays
    return _ObservationArrays(xyz, centers, normals, rows, weights, None)


def numba_continuous_time_identity(
    *, summary_parallel_cpu: bool = False, linearization_cache_cpu: bool = False,
    linearization_parallel_cpu: bool = False,
) -> str:
    """Identify the optional runtime and its arithmetic/compilation policy."""
    return (
        f"calibrex.numba_continuous_time/v0.9;numba={version('numba')};fastmath=false;cache=true;"
        f"summary_parallel_cpu={str(summary_parallel_cpu).lower()};summary_workers=3;"
        "summary_min_rows=32768;summary_workspace_bytes=16777216;"
        f"linearization_cache_cpu={str(linearization_cache_cpu).lower()};"
        "linearization_cache_bytes=67108864;"
        f"linearization_parallel_cpu={str(linearization_parallel_cpu).lower()};"
        "linearization_min_rows=32768;linearization_workspace_bytes=33554432"
    )


@njit(cache=True, nogil=True, fastmath=False, inline="always")  # type: ignore[untyped-decorator]
def _normalized(q: Quaternion) -> Quaternion:
    x, y, z, w = q
    norm = math.sqrt(x * x + y * y + z * z + w * w)
    return x / norm, y / norm, z / norm, w / norm


@njit(cache=True, nogil=True, fastmath=False, inline="always")  # type: ignore[untyped-decorator]
def _product(a: Quaternion, b: Quaternion) -> Quaternion:
    x, y, z, w = a
    rx, ry, rz, rw = b
    return (
        w * rx + x * rw + y * rz - z * ry,
        w * ry - x * rz + y * rw + z * rx,
        w * rz + x * ry - y * rx + z * rw,
        w * rw - x * rx - y * ry - z * rz,
    )


@njit(cache=True, nogil=True, fastmath=False, inline="always")  # type: ignore[untyped-decorator]
def _rotated(q: Quaternion, p: Vector3) -> Vector3:
    q = _normalized(q)
    x, y, z, w = _normalized(q)
    result = _product(_product(q, (p[0], p[1], p[2], 0.0)), (-x, -y, -z, w))
    return result[0], result[1], result[2]


@njit(cache=True, nogil=True, fastmath=False, inline="always")  # type: ignore[untyped-decorator]
def _composed(
    t: Vector3,
    q: Quaternion,
    rt: Vector3,
    rq: Quaternion,
) -> tuple[Vector3, Quaternion]:
    translated = _rotated(q, rt)
    return (
        (t[0] + translated[0], t[1] + translated[1], t[2] + translated[2]),
        _normalized(_normalized(_product(q, rq))),
    )


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _pose_world_rows(
    queries: IntArray,
    times: IntArray,
    samples: FloatArray,
    q0: FloatArray,
    q1: FloatArray,
    theta: FloatArray,
    sine: FloatArray,
    base: FloatArray,
    extrinsic: FloatArray,
    points: FloatArray,
    tolerance: float,
) -> tuple[FloatArray, FloatArray, NDArray[np.bool_], FloatArray, int, int]:
    rows = np.empty((len(queries), 7), dtype=np.float64)
    world = np.empty((len(queries), 3), dtype=np.float64)
    clamps = np.zeros(len(queries), dtype=np.bool_)
    extrapolation = np.zeros(len(queries), dtype=np.float64)
    calculated = 0
    bt = (base[0], base[1], base[2])
    bq = (base[3], base[4], base[5], base[6])
    et = (extrinsic[0], extrinsic[1], extrinsic[2])
    eq = (extrinsic[3], extrinsic[4], extrinsic[5], extrinsic[6])
    for row in range(len(queries)):
        timestamp = queries[row]
        if timestamp <= times[0] or timestamp >= times[-1]:
            index = 0 if timestamp <= times[0] else len(times) - 1
            clamps[row] = timestamp < times[0] or timestamp > times[-1]
            extrapolation[row] = max(
                0.0, float(times[0] - timestamp) / 1e9, float(timestamp - times[-1]) / 1e9
            )
            t = (samples[index, 0], samples[index, 1], samples[index, 2])
            q = (samples[index, 3], samples[index, 4], samples[index, 5], samples[index, 6])
        else:
            lo, hi = 0, len(times)
            while lo < hi:
                mid = (lo + hi) // 2
                if timestamp < times[mid]:
                    hi = mid
                else:
                    lo = mid + 1
            left, right = lo - 1, lo
            alpha = float(timestamp - times[left]) / float(times[right] - times[left])
            t = (
                samples[left, 0] + (samples[right, 0] - samples[left, 0]) * alpha,
                samples[left, 1] + (samples[right, 1] - samples[left, 1]) * alpha,
                samples[left, 2] + (samples[right, 2] - samples[left, 2]) * alpha,
            )
            if theta[left] < 0.0:
                q = _normalized(
                    (
                        q0[left, 0] + alpha * (q1[left, 0] - q0[left, 0]),
                        q0[left, 1] + alpha * (q1[left, 1] - q0[left, 1]),
                        q0[left, 2] + alpha * (q1[left, 2] - q0[left, 2]),
                        q0[left, 3] + alpha * (q1[left, 3] - q0[left, 3]),
                    )
                )
            else:
                wl = math.sin((1.0 - alpha) * theta[left]) / sine[left]
                wr = math.sin(alpha * theta[left]) / sine[left]
                q = _normalized(
                    (
                        wl * q0[left, 0] + wr * q1[left, 0],
                        wl * q0[left, 1] + wr * q1[left, 1],
                        wl * q0[left, 2] + wr * q1[left, 2],
                        wl * q0[left, 3] + wr * q1[left, 3],
                    )
                )
            q = _normalized(q)
            calculated += 1
        if extrapolation[row] > tolerance:
            return rows, world, clamps, extrapolation, calculated, row + 1
        t, q = _composed(t, q, bt, bq)
        rows[row] = (t[0], t[1], t[2], q[0], q[1], q[2], q[3])
        wt, wq = _composed(t, q, et, eq)
        p = _rotated(wq, (points[row, 0], points[row, 1], points[row, 2]))
        world[row] = (p[0] + wt[0], p[1] + wt[1], p[2] + wt[2])
    return rows, world, clamps, extrapolation, calculated, len(queries)


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _residual_rows(
    points: FloatArray,
    planes: FloatArray,
    normals: FloatArray,
    poses: FloatArray,
    weights: FloatArray,
    transform: FloatArray,
) -> FloatArray:
    result = np.empty(len(points), dtype=np.float64)
    t = (transform[0], transform[1], transform[2])
    q = (transform[3], transform[4], transform[5], transform[6])
    for row in range(len(points)):
        pt = (poses[row, 0], poses[row, 1], poses[row, 2])
        pq = (poses[row, 3], poses[row, 4], poses[row, 5], poses[row, 6])
        wt, wq = _composed(pt, pq, t, q)
        rotated = _rotated(wq, (points[row, 0], points[row, 1], points[row, 2]))
        dx = rotated[0] + wt[0] - planes[row, 0]
        dy = rotated[1] + wt[1] - planes[row, 1]
        dz = rotated[2] + wt[2] - planes[row, 2]
        result[row] = weights[row] * (
            normals[row, 0] * dx + normals[row, 1] * dy + normals[row, 2] * dz
        )
    return result


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _residual_rows_into(
    points: FloatArray, planes: FloatArray, normals: FloatArray, poses: FloatArray,
    weights: FloatArray, transform: FloatArray, result: FloatArray, start: int, stop: int,
) -> None:
    """Compute independent rows into a disjoint slice without changing arithmetic."""
    t = (transform[0], transform[1], transform[2])
    q = (transform[3], transform[4], transform[5], transform[6])
    for row in range(start, stop):
        pt = (poses[row, 0], poses[row, 1], poses[row, 2])
        pq = (poses[row, 3], poses[row, 4], poses[row, 5], poses[row, 6])
        wt, wq = _composed(pt, pq, t, q)
        rotated = _rotated(wq, (points[row, 0], points[row, 1], points[row, 2]))
        dx = rotated[0] + wt[0] - planes[row, 0]
        dy = rotated[1] + wt[1] - planes[row, 1]
        dz = rotated[2] + wt[2] - planes[row, 2]
        result[row] = weights[row] * (
            normals[row, 0] * dx + normals[row, 1] * dy + normals[row, 2] * dz
        )


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _world_rows(poses: FloatArray, extrinsic: FloatArray, points: FloatArray) -> FloatArray:
    world = np.empty((len(poses), 3), dtype=np.float64)
    et = (extrinsic[0], extrinsic[1], extrinsic[2])
    eq = (extrinsic[3], extrinsic[4], extrinsic[5], extrinsic[6])
    for row in range(len(poses)):
        t = (poses[row, 0], poses[row, 1], poses[row, 2])
        q = (poses[row, 3], poses[row, 4], poses[row, 5], poses[row, 6])
        wt, wq = _composed(t, q, et, eq)
        p = _rotated(wq, (points[row, 0], points[row, 1], points[row, 2]))
        world[row] = (p[0] + wt[0], p[1] + wt[1], p[2] + wt[2])
    return world


def numba_rig_residuals(arrays: _ObservationArrays, transform: SE3) -> list[float]:
    """Fuse the existing ordered residual formula without changing differences."""
    values = np.asarray((*transform.translation_m, *transform.rotation_quat_xyzw), dtype=np.float64)
    return list(
        _residual_rows(
            arrays.points,
            arrays.plane_points,
            arrays.normals,
            arrays.poses,
            arrays.sqrt_weights,
            values,
        ).tolist()
    )


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _jacobian_rows(
    points: FloatArray,
    planes: FloatArray,
    normals: FloatArray,
    poses: FloatArray,
    weights: FloatArray,
    transforms: FloatArray,
    steps: FloatArray,
) -> tuple[FloatArray, FloatArray]:
    residuals = _residual_rows(points, planes, normals, poses, weights, transforms[0])
    jacobian = np.empty((len(points), 6), dtype=np.float64)
    for column in range(6):
        plus = _residual_rows(points, planes, normals, poses, weights, transforms[1 + 2 * column])
        minus = _residual_rows(points, planes, normals, poses, weights, transforms[2 + 2 * column])
        for row in range(len(points)):
            jacobian[row, column] = (plus[row] - minus[row]) / (2.0 * steps[column])
    return residuals, jacobian


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _normal_rows(
    rows: FloatArray,
    residuals: FloatArray,
    huber: bool,
    delta: float,
) -> tuple[FloatArray, FloatArray]:
    hessian = np.zeros((6, 6), dtype=np.float64)
    gradient = np.zeros(6, dtype=np.float64)
    for index in range(len(rows)):
        absolute = abs(residuals[index])
        weight = delta / absolute if huber and absolute > delta and absolute != 0.0 else 1.0
        for left in range(6):
            weighted = weight * rows[index, left]
            gradient[left] += weighted * residuals[index]
            for right in range(6):
                hessian[left, right] += weighted * rows[index, right]
    return hessian, gradient


@njit(cache=True, nogil=True, fastmath=False, inline="never")  # type: ignore[untyped-decorator]
def _one_residual(
    point: FloatArray, plane: FloatArray, normal: FloatArray, pose: FloatArray,
    weight: float, transform: FloatArray,
) -> float:
    wt, wq = _composed(
        (pose[0], pose[1], pose[2]), (pose[3], pose[4], pose[5], pose[6]),
        (transform[0], transform[1], transform[2]),
        (transform[3], transform[4], transform[5], transform[6]),
    )
    rotated = _rotated(wq, (point[0], point[1], point[2]))
    dx, dy, dz = (rotated[0]+wt[0]-plane[0], rotated[1]+wt[1]-plane[1],
                  rotated[2]+wt[2]-plane[2])
    return float(weight * (normal[0]*dx + normal[1]*dy + normal[2]*dz))


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _residual_square_sum(
    points: FloatArray, planes: FloatArray, normals: FloatArray, poses: FloatArray,
    weights: FloatArray, transform: FloatArray, compensated: bool,
) -> float:
    total, correction = 0.0, 0.0
    for row in range(len(points)):
        value = _one_residual(points[row], planes[row], normals[row], poses[row],
                              weights[row], transform)
        square = value * value
        updated = total + square
        # Match CPython's ordered float sum: plain in 3.10/3.11, Neumaier in 3.12.
        if compensated:
            if abs(total) >= abs(square):
                correction += (total-updated) + square
            else:
                correction += (square-updated) + total
        total = updated
    if compensated and correction != 0.0 and math.isfinite(correction):
        total += correction
    return total


def _numba_rmse_policy() -> bool | None:
    """Select only the interpreter summation policies whose bits are validated."""
    if sys.implementation.name != "cpython" or sys.version_info[:2] not in (
        (3, 10), (3, 11), (3, 12),
    ):
        return None
    return sys.version_info[:2] == (3, 12)


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _value_square_sum(values: FloatArray, compensated: bool) -> float:
    total, correction = 0.0, 0.0
    for value in values:
        square = value * value
        updated = total + square
        if compensated:
            if abs(total) >= abs(square):
                correction += (total-updated) + square
            else:
                correction += (square-updated) + total
        total = updated
    if compensated and correction != 0.0 and math.isfinite(correction):
        total += correction
    return total


def _numba_values_rmse(values: FloatArray) -> tuple[float] | None:
    """Reduce captured native residuals in their exact Python summation order."""
    policy = _numba_rmse_policy()
    if policy is None or not len(values):
        return None
    total = _value_square_sum(values, policy)
    return (math.sqrt(total / len(values)),)


def _numba_rig_rmse(arrays: _ObservationArrays, transform: SE3) -> tuple[float] | None:
    """Compute native residual RMSE without creating row arrays or Python floats."""
    policy = _numba_rmse_policy()
    if policy is None:
        return None
    values = np.asarray((*transform.translation_m, *transform.rotation_quat_xyzw), dtype=np.float64)
    square_sum = _residual_square_sum(
        arrays.points, arrays.plane_points, arrays.normals, arrays.poses,
        arrays.sqrt_weights, values, policy,
    )
    return (math.sqrt(square_sum / len(arrays.points)),) if len(arrays.points) else None


def _numba_parallel_rig_rmse(
    arrays: _ObservationArrays, transform: SE3, executor: ThreadPoolExecutor,
) -> tuple[float] | None:
    """Split row computation, then retain the exact ordered CPU reduction."""
    n = len(arrays.points)
    policy = _numba_rmse_policy()
    if (policy is None or n < 32768 or n * 8 > 16 * 1024**2
        or any(value.flags.writeable for value in (
            arrays.points, arrays.plane_points, arrays.normals, arrays.poses, arrays.sqrt_weights,
        ))):
        return None
    values = np.asarray((*transform.translation_m, *transform.rotation_quat_xyzw), dtype=np.float64)
    output = np.empty(n, dtype=np.float64)
    midpoint = n // 2
    arguments = (arrays.points, arrays.plane_points, arrays.normals, arrays.poses,
                 arrays.sqrt_weights, values, output)
    future = executor.submit(
        _residual_rows_into, arrays.points, arrays.plane_points, arrays.normals,
        arrays.poses, arrays.sqrt_weights, values, output, midpoint, n,
    )
    try:
        _residual_rows_into(*arguments, 0, midpoint)
    except BaseException:
        # No task may outlive this request's output, including on an exception.
        with suppress(BaseException):
            future.result()
        raise
    future.result()
    total = _value_square_sum(output, policy)
    return (math.sqrt(total / n),)


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _fused_jacobian_rows(
    points: FloatArray, planes: FloatArray, normals: FloatArray, poses: FloatArray,
    weights: FloatArray, transforms: FloatArray, steps: FloatArray,
) -> tuple[FloatArray, FloatArray]:
    residuals = np.empty(len(points), dtype=np.float64)
    jacobian = np.empty((len(points), 6), dtype=np.float64)
    for row in range(len(points)):
        point, plane, normal, pose, weight = (
            points[row], planes[row], normals[row], poses[row], weights[row],
        )
        residuals[row] = _one_residual(point, plane, normal, pose, weight, transforms[0])
        for column in range(6):
            plus = _one_residual(point, plane, normal, pose, weight, transforms[1+2*column])
            minus = _one_residual(point, plane, normal, pose, weight, transforms[2+2*column])
            jacobian[row, column] = (plus-minus)/(2.0*steps[column])
    return residuals, jacobian


def numba_rig_linearization(
    arrays: _ObservationArrays,
    transforms: Sequence[SE3],
    steps: Sequence[float],
) -> tuple[FloatArray, FloatArray]:
    """Keep all thirteen finite-difference evaluations in owned numeric arrays."""
    values = np.asarray(
        [(*t.translation_m, *t.rotation_quat_xyzw) for t in transforms], dtype=np.float64
    )
    arguments = (arrays.points, arrays.plane_points, arrays.normals, arrays.poses,
                 arrays.sqrt_weights, values, np.asarray(steps, dtype=np.float64))
    if len(arrays.points) >= 4096:
        if any(step == 0.0 for step in steps):
            raise ZeroDivisionError("division by zero")
        return cast(tuple[FloatArray, FloatArray], _fused_jacobian_rows(*arguments))
    return cast(
        tuple[FloatArray, FloatArray],
        _jacobian_rows(*arguments),
    )


@njit(cache=True, nogil=True, fastmath=False)  # type: ignore[untyped-decorator]
def _jacobian_rows_into(
    points: FloatArray, planes: FloatArray, normals: FloatArray, poses: FloatArray,
    weights: FloatArray, transforms: FloatArray, steps: FloatArray,
    residuals: FloatArray, jacobian: FloatArray, plus: FloatArray, minus: FloatArray,
    start: int, stop: int,
) -> None:
    """Reuse two scratch columns within each worker's disjoint row range."""
    _residual_rows_into(points, planes, normals, poses, weights, transforms[0],
                        residuals, start, stop)
    for column in range(6):
        _residual_rows_into(points, planes, normals, poses, weights, transforms[1+2*column],
                            plus, start, stop)
        _residual_rows_into(points, planes, normals, poses, weights, transforms[2+2*column],
                            minus, start, stop)
        for row in range(start, stop):
            jacobian[row, column] = (plus[row]-minus[row])/(2.0*steps[column])


def _numba_parallel_rig_linearization(
    arrays: _ObservationArrays, values: FloatArray, steps: FloatArray,
    executor: ThreadPoolExecutor,
) -> tuple[FloatArray, FloatArray] | None:
    """Compute independent differences; leave all reductions in their original order."""
    n = len(arrays.points)
    if (n < 32768 or n * 16 > 32 * 1024**2
        or any(value.flags.writeable for value in (
            arrays.points, arrays.plane_points, arrays.normals, arrays.poses, arrays.sqrt_weights,
        ))):
        return None
    residuals = np.empty(n, dtype=np.float64)
    jacobian = np.empty((n, 6), dtype=np.float64)
    plus, minus = np.empty(n, dtype=np.float64), np.empty(n, dtype=np.float64)
    arguments = (arrays.points, arrays.plane_points, arrays.normals, arrays.poses,
                 arrays.sqrt_weights, values, steps, residuals, jacobian, plus, minus)
    midpoint = n // 2
    future = executor.submit(
        _jacobian_rows_into, arrays.points, arrays.plane_points, arrays.normals,
        arrays.poses, arrays.sqrt_weights, values, steps, residuals, jacobian, plus, minus,
        midpoint, n,
    )
    try:
        _jacobian_rows_into(*arguments, 0, midpoint)
    except BaseException:
        with suppress(BaseException):
            future.result()
        raise
    future.result()
    return residuals, jacobian


class NumbaRigLinearizer:
    """Own the last immutable CPU difference batch with a bounded host cache."""

    def __init__(self, executor: ThreadPoolExecutor | None = None) -> None:
        self._executor = executor
        self._arrays: _ObservationArrays | None = None
        self._key: tuple[bytes, bytes] | None = None
        self._result: tuple[FloatArray, FloatArray] | None = None
        self._lock = Lock()
        self.calls = 0
        self.parallel_calls = 0
        self.cache_hits = 0
        self.cache_bytes = 0
        self.rows = 0

    def _clear(self) -> None:
        self._arrays = None
        self._key = None
        self._result = None
        self.cache_bytes = 0

    def linearize(
        self, arrays: _ObservationArrays, transforms: Sequence[SE3], steps: Sequence[float],
    ) -> tuple[FloatArray, FloatArray]:
        """Retain exact differences while giving every caller writable owned arrays."""
        if len(transforms) != 13 or len(steps) != 6:
            raise ValueError("CPU rig differences require thirteen transforms and six steps")
        if any(step == 0.0 for step in steps):
            raise ZeroDivisionError("division by zero")
        values = np.asarray(
            [(*t.translation_m, *t.rotation_quat_xyzw) for t in transforms], dtype=np.float64,
        )
        key = values.tobytes(), np.asarray(steps, dtype=np.float64).tobytes()
        immutable = not any(value.flags.writeable for value in (
            arrays.points, arrays.plane_points, arrays.normals, arrays.poses, arrays.sqrt_weights,
        ))
        with self._lock:
            if (immutable and self._arrays is arrays and self._key == key
                and self._result is not None):
                self.cache_hits += 1
                return self._result[0].copy(), self._result[1].copy()
            self._clear()
            result = (_numba_parallel_rig_linearization(
                arrays, values, np.asarray(steps, dtype=np.float64), self._executor,
            ) if self._executor is not None else None)
            if result is None:
                result = numba_rig_linearization(arrays, transforms, steps)
            else:
                self.parallel_calls += 1
            residuals, jacobian = result
            self.calls += 1
            self.rows += len(arrays.points)
            size = residuals.nbytes + jacobian.nbytes
            if immutable and size <= 64 * 1024**2:
                saved_r, saved_j = residuals.copy(), jacobian.copy()
                saved_r.flags.writeable = saved_j.flags.writeable = False
                self._arrays, self._key, self._result = arrays, key, (saved_r, saved_j)
                self.cache_bytes = size
            return residuals, jacobian


def numba_rig_normal_equations(
    rows: FloatArray,
    residuals: FloatArray,
    huber: bool,
    delta: float,
) -> tuple[list[list[float]], list[float]]:
    """Accumulate every row/product in the original robust weighted order."""
    hessian, gradient = _normal_rows(rows, residuals, huber, delta)
    return hessian.tolist(), gradient.tolist()


@dataclass(frozen=True, eq=False)
class _CapturedTimestampRows:
    """Private immutable native capture times, validated once per selected split."""

    values: IntArray
    minimum: int
    maximum: int

    @classmethod
    def capture(cls, timestamps: Sequence[int]) -> _CapturedTimestampRows | None:
        if (type(timestamps) not in (tuple, list) or not timestamps
            or any(type(t) is not int or abs(t) >= 2**63 for t in timestamps)):
            return None
        values = np.asarray(timestamps, dtype=np.int64)
        values.flags.writeable = False
        return cls(values, int(values.min()), int(values.max()))

    def shifted(self, offset_ns: int, origin_ns: int) -> _PreparedTimestampQuery | None:
        if type(offset_ns) is not int or abs(offset_ns) >= 2**63 or type(origin_ns) is not int:
            return None
        lo, hi = self.minimum+offset_ns, self.maximum+offset_ns
        if (lo <= -(2**63) or hi >= 2**63
            or max(abs(lo-origin_ns), abs(hi-origin_ns)) > 2**53):
            return None
        values = self.values+np.int64(offset_ns)
        values.flags.writeable = False
        return _PreparedTimestampQuery(self, offset_ns, origin_ns, values)


@dataclass(frozen=True, eq=False)
class _PreparedTimestampQuery:
    owner: _CapturedTimestampRows
    offset_ns: int
    origin_ns: int
    values: IntArray

    def __hash__(self) -> int:
        return hash((self.owner, self.offset_ns, self.origin_ns))

    def __eq__(self, other: object) -> bool:
        return (isinstance(other, _PreparedTimestampQuery) and self.owner is other.owner
                and self.offset_ns == other.offset_ns and self.origin_ns == other.origin_ns)


_PoseKey: TypeAlias = tuple[tuple[int, ...] | _PreparedTimestampQuery, bytes]
_NATIVE_RECORD_EXTRAPOLATION = OdometryTrack._record_extrapolation


@dataclass(frozen=True)
class _PreparedPoseBatch:
    poses: tuple[SE3, ...]
    rows: FloatArray
    clamps: int
    max_extrapolation_s: float
    size_bytes: int


@dataclass
class NumbaContinuousTimeQueryContext:
    """Capture a supported immutable trajectory once per solve."""

    times: IntArray
    samples: FloatArray
    q0: FloatArray
    q1: FloatArray
    theta: FloatArray
    sine: FloatArray
    sample_snapshot: object
    timestamp_snapshot: object
    used: bool = False
    _pose_cache_limit_bytes: int = 128 * 1024 * 1024
    _pose_cache_bytes: int = 0
    _pose_cache: OrderedDict[_PoseKey, _PreparedPoseBatch] = field(default_factory=OrderedDict)

    @classmethod
    def capture(cls, track: OdometryTrack) -> NumbaContinuousTimeQueryContext | None:
        """Return None for trajectories requiring the original scalar semantics."""
        if (
            type(track) is not OdometryTrack
            or "interpolate" in track.__dict__
            or len(track.samples) < 2
        ):
            return None
        if any(
            type(sample.pose) is not SE3 or type(sample.timestamp_ns) is not int
            for sample in track.samples
        ):
            return None
        raw_times = track._timestamps_ns
        if (
            raw_times != tuple(sample.timestamp_ns for sample in track.samples)
            or any(abs(timestamp) >= 2**63 for timestamp in raw_times)
            or raw_times[-1] - raw_times[0] > 2**53
            or any(a >= b for a, b in pairwise(raw_times))
        ):
            return None
        times = np.asarray(raw_times, dtype=np.int64)
        samples = np.asarray(
            [
                (*sample.pose.translation_m, *sample.pose.rotation_quat_xyzw)
                for sample in track.samples
            ],
            dtype=np.float64,
        )
        if not np.isfinite(samples).all():
            return None
        prepared = [_PreparedInterval.capture(a.pose, b.pose) for a, b in pairwise(track.samples)]
        q0 = np.asarray([interval.q0 for interval in prepared], dtype=np.float64)
        q1 = np.asarray([interval.q1 for interval in prepared], dtype=np.float64)
        theta = np.asarray(
            [-1.0 if interval.theta is None else interval.theta for interval in prepared],
            dtype=np.float64,
        )
        sine = np.asarray([interval.sin_theta for interval in prepared], dtype=np.float64)
        for array in (times, samples, q0, q1, theta, sine):
            array.flags.writeable = False
        return cls(times, samples, q0, q1, theta, sine, track.samples, track._timestamps_ns)

    def evaluate(
        self,
        track: OdometryTrack,
        timestamps: list[int],
        points: list[Vector3],
        base: SE3,
        extrinsic: SE3,
        tolerance: float,
    ) -> tuple[list[SE3], FloatArray, bool] | None:
        result = self._evaluate(track, timestamps, points, base, extrinsic, tolerance, False)
        return None if result is None else (result[0], result[2], result[3])

    def _evaluate_rows(
        self, track: OdometryTrack, timestamps: list[int], points: FloatArray,
        base: SE3, extrinsic: SE3, tolerance: float,
    ) -> tuple[FloatArray, FloatArray, bool] | None:
        result = self._evaluate(track, timestamps, points, base, extrinsic, tolerance, True)
        return None if result is None else (result[1], result[2], result[3])

    def _evaluate_prepared_rows(
        self, track: OdometryTrack, timestamps: _PreparedTimestampQuery, points: FloatArray,
        base: SE3, extrinsic: SE3, tolerance: float,
    ) -> tuple[FloatArray, FloatArray, bool] | None:
        result = self._evaluate(track, timestamps, points, base, extrinsic, tolerance, True)
        return None if result is None else (result[1], result[2], result[3])

    def _evaluate(
        self, track: OdometryTrack, timestamps: list[int] | _PreparedTimestampQuery,
        points: list[Vector3] | FloatArray,
        base: SE3, extrinsic: SE3, tolerance: float, rows_only: bool,
    ) -> tuple[list[SE3], FloatArray, FloatArray, bool] | None:
        """Preserve scalar query counters and stop at the first extrapolation failure."""
        if (
            track.samples is not self.sample_snapshot
            or track._timestamps_ns is not self.timestamp_snapshot
            or "interpolate" in track.__dict__
            or type(base) is not SE3
            or type(extrinsic) is not SE3
        ):
            return None
        if isinstance(timestamps, _PreparedTimestampQuery):
            if timestamps.origin_ns != int(self.times[0]):
                return None
            query_values = timestamps.values
            key_times: tuple[int, ...] | _PreparedTimestampQuery = timestamps
        else:
            if any(type(t) is not int or abs(t) >= 2**63
                   or abs(t-int(self.times[0])) > 2**53 for t in timestamps):
                return None
            query_values = np.asarray(timestamps, dtype=np.int64)
            key_times = tuple(timestamps)
        base_values = np.asarray((*base.translation_m, *base.rotation_quat_xyzw), dtype=np.float64)
        transform = np.asarray((*extrinsic.translation_m, *extrinsic.rotation_quat_xyzw))
        coordinates = np.asarray(points, dtype=np.float64).reshape((-1, 3))
        if not all(np.isfinite(array).all() for array in (base_values, transform, coordinates)):
            return None
        key = (key_times, base_values.tobytes())
        can_reuse = (
            self._pose_cache_limit_bytes > 0
            and "_record_extrapolation" not in track.__dict__
            and OdometryTrack._record_extrapolation is _NATIVE_RECORD_EXTRAPOLATION
        )
        with track._cache_lock:
            cached = self._pose_cache.get(key) if can_reuse else None
            if cached is not None and cached.max_extrapolation_s <= tolerance:
                self._pose_cache.move_to_end(key)
                track.interpolation_count += len(query_values)
                track.clamp_count += cached.clamps
                track._record_extrapolation(cached.max_extrapolation_s)
            else:
                cached = None
        if cached is not None:
            # Only trajectory/base geometry is retained. Changed points and
            # extrinsics always get the same ordered world-coordinate formula.
            world = _world_rows(cached.rows, transform, coordinates)
            self.used = True
            poses = [] if rows_only else (
                list(cached.poses) if cached.poses else _materialize_pose_rows(cached.rows)
            )
            return poses, cached.rows, world, True
        rows, world, clamps, extrapolation, _calculated, visited = _pose_world_rows(
            query_values,
            self.times,
            self.samples,
            self.q0,
            self.q1,
            self.theta,
            self.sine,
            base_values,
            transform,
            coordinates,
            tolerance,
        )
        self.used = True
        with track._cache_lock:
            track.interpolation_count += visited
            track.clamp_count += int(np.count_nonzero(clamps[:visited]))
            if ("_record_extrapolation" not in track.__dict__
                and OdometryTrack._record_extrapolation is _NATIVE_RECORD_EXTRAPOLATION):
                if visited:
                    track._record_extrapolation(float(np.max(extrapolation[:visited])))
            else:
                for value in extrapolation[:visited]:
                    track._record_extrapolation(float(value))
        usable = visited == len(query_values) and not np.any(extrapolation[:visited] > tolerance)
        # These rows have already performed both public-constructor normalizations.
        poses = []
        if usable:
            if not rows_only:
                poses = _materialize_pose_rows(rows)
            if can_reuse and len(rows):
                pose_tuple = tuple(poses)
                # Charge all retained payloads, even shared Python floats/ints,
                # plus a conservative allowance for the LRU node and entry.
                size = sys.getsizeof(rows) + sys.getsizeof(pose_tuple) + sys.getsizeof(key)
                # Numba arrays retain an NRT owner, so their Python header does
                # not include the backing allocation in sys.getsizeof.
                if not rows.flags.owndata:
                    size += rows.nbytes
                size += sys.getsizeof(key[0]) + sys.getsizeof(key[1]) + 512
                if isinstance(key[0], _PreparedTimestampQuery):
                    size += sys.getsizeof(key[0].values) + sys.getsizeof(key[0].owner)
                    size += sys.getsizeof(key[0].owner.values) + 128
                    size += sys.getsizeof(key[0].__dict__) + sys.getsizeof(key[0].owner.__dict__)
                else:
                    size += len(key[0]) * sys.getsizeof(2**63 - 1)
                if poses:
                    prototype = poses[0]
                    pose_size = (
                        sys.getsizeof(prototype)
                        + max(sys.getsizeof(pose.__dict__) for pose in poses)
                        + sys.getsizeof(prototype.translation_m)
                        + sys.getsizeof(prototype.rotation_quat_xyzw) + 7 * sys.getsizeof(0.0)
                    )
                    size += len(poses) * pose_size
                if size <= self._pose_cache_limit_bytes:
                    rows.flags.writeable = False
                    batch = _PreparedPoseBatch(
                        pose_tuple, rows, int(np.count_nonzero(clamps)),
                        float(np.max(extrapolation)), size,
                    )
                    with track._cache_lock:
                        previous = self._pose_cache.pop(key, None)
                        if previous is not None:
                            self._pose_cache_bytes -= previous.size_bytes
                        while self._pose_cache and self._pose_cache_bytes + size > (
                            self._pose_cache_limit_bytes
                        ):
                            evicted = self._pose_cache.popitem(last=False)[1]
                            self._pose_cache_bytes -= evicted.size_bytes
                        self._pose_cache[key] = batch
                        self._pose_cache_bytes += size
        rows.flags.writeable = False
        return poses, rows, world, usable


def _materialize_pose_rows(rows: FloatArray) -> list[SE3]:
    poses = []
    for start in range(0, len(rows), 4096):
        for row in rows[start : start + 4096].tolist():
            pose = object.__new__(SE3)
            object.__setattr__(pose, "translation_m", tuple(row[:3]))
            object.__setattr__(pose, "rotation_quat_xyzw", tuple(row[3:]))
            poses.append(pose)
    return poses
