"""Optional device-resident point-to-plane registration for LiDAR odometry."""

from __future__ import annotations

import hashlib
import math
from dataclasses import dataclass, replace
from typing import Any, Protocol

import cupy as cp
import numpy as np
from numpy.typing import NDArray
from scipy.spatial import cKDTree

from calibrex.solvers.cuda_lidar_map_adapter import (
    _SOURCE,
    _coordinates,
    _CudaTree,
    _host,
)
from calibrex.solvers.scan_to_scan_odometry import (
    FloatArray,
    ScanOdometryOptions,
    ScanRegistration,
    _canonicalize_motion,
    _exp_se3,
    _norm3,
    _normals_from_neighbours,
    _quantize_geometry,
    register_point_to_plane,
)

_SYSTEM_SOURCE = r"""
extern "C" __global__ void system3(
    const double* moved, const double* target, const double* normal,
    const bool* planar, const double* distances, const long long* indices,
    const int n, const int m, const double gate, const double huber,
    double* terms) {
    int row=blockDim.x*blockIdx.x+threadIdx.x;
    if(row>=n) return;
    long long hit=indices[2*row];
    bool valid=distances[2*row]<gate && hit<m;
    if(valid) valid=planar[hit];
    double j[6]={0,0,0,0,0,0}, r=0, w=1;
    if(valid){
        const double* p=moved+3*row; const double* q=target+3*hit;
        const double* a=normal+3*hit;
        r=(p[0]-q[0])*a[0]+(p[1]-q[1])*a[1]+(p[2]-q[2])*a[2];
        j[0]=p[1]*a[2]-p[2]*a[1];j[1]=p[2]*a[0]-p[0]*a[2];
        j[2]=p[0]*a[1]-p[1]*a[0];j[3]=a[0];j[4]=a[1];j[5]=a[2];
        w=fabs(r)<=huber?1.0:huber/fmax(fabs(r),1e-12);
    }
    int column=0;
    for(int a=0;a<6;a++) for(int b=a;b<6;b++)
        terms[(column++)*n+row]=(j[a]*w)*j[b];
    for(int a=0;a<6;a++) terms[(column++)*n+row]=(j[a]*w)*r;
    terms[27*n+row]=r*r; terms[28*n+row]=valid?1.0:0.0;
}
"""


_COVARIANCE_SOURCE = r"""
extern "C" __global__ void covariance3(
    const double* points, const long long* indices, const int n, const int k,
    double* covariance) {
    int row=blockDim.x*blockIdx.x+threadIdx.x;
    if(row>=n) return;
    double mean[3]={0,0,0};
    for(int j=0;j<k;j++) for(int a=0;a<3;a++) mean[a]+=points[3*indices[row*(k+1)+j]+a];
    for(int a=0;a<3;a++) mean[a]/=k;
    double c[9]={0,0,0,0,0,0,0,0,0};
    for(int j=0;j<k;j++){
        double d[3];for(int a=0;a<3;a++) d[a]=points[3*indices[row*(k+1)+j]+a]-mean[a];
        for(int a=0;a<3;a++) for(int b=0;b<3;b++) c[3*a+b]+=d[a]*d[b];
    }
    for(int j=0;j<9;j++) covariance[9*row+j]=c[j]/k;
}
"""

_NORMAL_SOURCE = r"""
// FP64 symmetric 3x3 PCA; uncertain memberships/eigenspaces retain native repairs.
extern "C" __global__ void normals3(
    const double* points, const long long* indices, const double* distances,
    const long long* rows, const int n, const int k, const double ratio,
    double* normal, bool* planar, bool* done, unsigned char* repair) {
    int row=blockDim.x*blockIdx.x+threadIdx.x;
    if(row>=n) return;
    const double eps=2.2204460492503131e-16;
    double mean[3]={0,0,0},c[9]={0,0,0,0,0,0,0,0,0};
    for(int j=0;j<k;j++) for(int a=0;a<3;a++)
        mean[a]+=points[3*indices[row*(k+1)+j]+a];
    for(int a=0;a<3;a++) mean[a]/=k;
    for(int j=0;j<k;j++) {
        double d[3];
        for(int a=0;a<3;a++) d[a]=points[3*indices[row*(k+1)+j]+a]-mean[a];
        for(int a=0;a<3;a++) for(int b=0;b<3;b++) c[3*a+b]+=d[a]*d[b];
    }
    for(int a=0;a<9;a++) c[a]/=k;
    double scale=fmax(fmax(fabs(c[0]),fabs(c[4])),fabs(c[8]));
    bool bad=!(scale>0) || !isfinite(scale);
    double unit=fmax(scale,2.2250738585072014e-308);
    for(int a=0;a<9;a++) c[a]/=unit;
    double trace=c[0]+c[4]+c[8],q=trace/3;
    double x=c[0]-q,y=c[4]-q,z=c[8]-q;
    double p=sqrt((x*x+y*y+z*z+2*(c[1]*c[1]+c[2]*c[2]+c[5]*c[5]))/6);
    bad|=!(p>64*eps);
    double inv=1/fmax(p,2.2250738585072014e-308);
    double b0=x*inv,b1=y*inv,b2=z*inv;
    double b3=c[1]*inv,b4=c[2]*inv,b5=c[5]*inv;
    double determinant=b0*b1*b2+2*b3*b4*b5-b0*b5*b5-b1*b4*b4-b2*b3*b3;
    double phi=acos(fmax(-1.0,fmin(1.0,determinant/2)))/3;
    double low=q+2*p*cos(phi+2.0943951023931954923);
    double high=q+2*p*cos(phi),mid=trace-low-high;
    double a[9];for(int j=0;j<9;j++) a[j]=c[j];
    a[0]-=low;a[4]-=low;a[8]-=low;
    double best=-1,v[3]={0,0,0};
    for(int i=0;i<3;i++) {
        int j=(i+1)%3;
        double w[3]={a[3*i+1]*a[3*j+2]-a[3*i+2]*a[3*j+1],
                     a[3*i+2]*a[3*j]-a[3*i]*a[3*j+2],
                     a[3*i]*a[3*j+1]-a[3*i+1]*a[3*j]};
        double norm=w[0]*w[0]+w[1]*w[1]+w[2]*w[2];
        if(norm>best) {best=norm;for(int u=0;u<3;u++) v[u]=w[u];}
    }
    bad|=!(best>0) || !isfinite(best);
    double length=sqrt(fmax(best,2.2250738585072014e-308));
    for(int u=0;u<3;u++) v[u]/=length;
    double residual=0;
    for(int i=0;i<3;i++) {
        double r=c[3*i]*v[0]+c[3*i+1]*v[1]+c[3*i+2]*v[2]-low*v[i];
        residual+=r*r;
    }
    double floor=1e-12/unit,spectral=fmax(fabs(high),2.2250738585072014e-308);
    double threshold=ratio*fmax(mid,floor);
    // Cover closed-form root error near repeated upper eigenvalues.
    bad|=mid-low<=1e-4*spectral || fabs(low-threshold)<=2e-8*spectral;
    bad|=!isfinite(residual) || sqrt(residual)>64*eps*spectral;
    bool tied=distances[row*(k+1)+k]-distances[row*(k+1)+k-1]
              <=64*eps*fmax(1.0,distances[row*(k+1)+k]);
    long long target=rows[row];
    for(int u=0;u<3;u++) normal[3*target+u]=v[u];
    planar[target]=low<=threshold;
    done[target]=true;
    repair[row]=(bad||tied)?1:0;
}
"""


_FUSED_SOURCE = r"""
extern "C" __global__ void project_query_system3(
    const double* xyz, const double* matrix, const double* reference,
    const long long* order, const int* axis, const double* split,
    const int* child, const long long* span, const double* boxes, const double* normal,
    const bool* planar, const bool* done, const int n, const int m, const double gate,
    const double huber, const double error, double* uncertain, double* terms,
    double* moved_out, double* distances_out, long long* indices_out,
    long long* pending_out) {
    int row=blockDim.x*blockIdx.x+threadIdx.x;
    if(row>=n) return;
    const double* s=xyz+3*row;
    double p[3];
    for(int a=0;a<3;a++)
        p[a]=matrix[4*a]*s[0]+matrix[4*a+1]*s[1]+matrix[4*a+2]*s[2]+matrix[4*a+3];
    double best[2]={1.0/0.0,1.0/0.0};long long chosen[2]={m,m};
    int stack[64];int top=1;stack[0]=0;
    while(top){
        --top;int node=stack[top];
        double d[3];
        for(int a=0;a<3;a++)
            d[a]=fmax(0.0,fmax(boxes[6*node+a]-p[a],p[a]-boxes[6*node+3+a]));
        double lower=(d[0]*d[0]+d[1]*d[1]+d[2]*d[2])*(1.0-16.0*2.2204460492503131e-16);
        if(lower>best[1]) continue;
        if(axis[node]<0){
            for(long long v=span[2*node];v<span[2*node+1];v++){
                long long position=order[v];
                double x=p[0]-reference[3*position];
                double y=p[1]-reference[3*position+1];
                double z=p[2]-reference[3*position+2];
                double d=x*x+y*y+z*z;
                if(d<best[1] || (d==best[1] && position<chosen[1])){
                    if(d<best[0] || (d==best[0] && position<chosen[0])){
                        best[1]=best[0];chosen[1]=chosen[0];best[0]=d;chosen[0]=position;
                    }else{best[1]=d;chosen[1]=position;}
                }
            }
        }else{
            double delta=p[axis[node]]-split[node];int side=delta<0?0:1;
            if(delta*delta<=best[1]){
                stack[top]=child[2*node+1-side];++top;
            }
            stack[top]=child[2*node+side];++top;
        }
    }
    double distance=sqrt(best[0]);long long hit=chosen[0];
    for(int a=0;a<3;a++) moved_out[3*row+a]=p[a];
    distances_out[2*row]=distance;distances_out[2*row+1]=sqrt(best[1]);
    indices_out[2*row]=hit;indices_out[2*row+1]=chosen[1];
    bool repair=sqrt(best[1])-distance<=error || fabs(distance-gate)<=error;
    uncertain[row]=repair?1.0:0.0;
    bool valid=distance<gate && hit<m;
    bool pending=valid && !done[hit];
    pending_out[row]=pending?hit:-1;
    if(valid) valid=done[hit] && planar[hit];
    double j[6]={0,0,0,0,0,0},r=0,w=1;
    if(valid){
        const double* q=reference+3*hit;const double* a=normal+3*hit;
        r=(p[0]-q[0])*a[0]+(p[1]-q[1])*a[1]+(p[2]-q[2])*a[2];
        j[0]=p[1]*a[2]-p[2]*a[1];j[1]=p[2]*a[0]-p[0]*a[2];
        j[2]=p[0]*a[1]-p[1]*a[0];j[3]=a[0];j[4]=a[1];j[5]=a[2];
        w=fabs(r)<=huber?1.0:huber/fmax(fabs(r),1e-12);
    }
    int column=0;
    for(int a=0;a<6;a++) for(int b=a;b<6;b++)
        terms[(column++)*n+row]=(j[a]*w)*j[b];
    for(int a=0;a<6;a++) terms[(column++)*n+row]=(j[a]*w)*r;
    terms[27*n+row]=r*r;terms[28*n+row]=valid?1.0:0.0;
    terms[29*n+row]=repair?1.0:0.0;
    terms[30*n+row]=pending?1.0:0.0;
}
"""


def execution_identity() -> str:
    """Identify arithmetic, query code and runtime without the device model."""
    properties = cp.cuda.runtime.getDeviceProperties(cp.cuda.runtime.getDevice())
    system_digest = hashlib.sha256(
        (_SYSTEM_SOURCE + _FUSED_SOURCE + _COVARIANCE_SOURCE + _NORMAL_SOURCE).encode(),
    ).hexdigest()
    return (
        "calibrex.cuda_scan_odometry/v0.1;fp64=true;query_fma=false;"
        f"cupy={cp.__version__};cuda={cp.cuda.runtime.runtimeGetVersion()};"
        f"driver={cp.cuda.runtime.driverGetVersion()};"
        f"cc={properties['major']}.{properties['minor']};"
        f"query_sha256={hashlib.sha256(_SOURCE.encode()).hexdigest()};"
        f"system_sha256={system_digest}"
        ";normal_schedule=async-v1"
    )


@dataclass
class _PendingNormals:
    rows: cp.ndarray
    flags: cp.ndarray
    distances: cp.ndarray
    indexes: cp.ndarray
    native_flags: NDArray[np.uint8]
    pinned: Any  # Own the pinned allocation until its download event completes.
    ready: cp.cuda.Event
    k: int


@dataclass
class _CapturedSource:
    """Own one register call's source; host copies are created only by native repairs."""

    host: FloatArray | None
    points: cp.ndarray | None = None

    @property
    def nbytes(self) -> int:
        if self.host is not None:
            return int(self.host.nbytes)
        assert self.points is not None
        return int(self.points.nbytes)

    def maximum(self) -> float:
        if self.host is not None:
            return float(np.max(np.abs(self.host))) if len(self.host) else 1.0
        assert self.points is not None
        with cp.cuda.Device(self.points.device.id):
            return float(cp.max(cp.abs(self.points)).get()) if len(self.points) else 1.0

    def device(self) -> cp.ndarray:
        if self.points is None:
            assert self.host is not None
            self.points = cp.asarray(self.host)
        return self.points

    def native(self) -> FloatArray:
        if self.host is None:
            assert self.points is not None
            with cp.cuda.Device(self.points.device.id):
                self.host = np.asarray(cp.asnumpy(self.points), dtype=np.float64)
        return self.host

    def native_rows(self, rows: NDArray[np.int64]) -> FloatArray:
        if self.host is not None:
            return self.host[rows]
        assert self.points is not None
        with cp.cuda.Device(self.points.device.id):
            return np.asarray(cp.asnumpy(self.points[cp.asarray(rows)]), dtype=np.float64)


class _TargetTree(Protocol):
    """Spatial index contract shared by captured and device-built targets."""

    data: cp.ndarray
    arrays: tuple[cp.ndarray, ...]
    capture_method: str
    query_calls: int
    query_rows: int

    @property
    def cpu(self) -> cKDTree: ...

    def query(self, points: cp.ndarray, k: int) -> tuple[cp.ndarray, cp.ndarray]: ...


def _validate_options(options: ScanOdometryOptions) -> None:
    if not 2 <= options.normal_neighbours <= 16:
        raise ValueError("CUDA normal neighborhoods require 2 through 16 neighbors")
    if not np.isfinite(options.max_correspondence_m) or options.max_correspondence_m <= 0:
        raise ValueError("CUDA correspondence gate must be finite and positive")


class CudaOdometryTarget:
    """Capture one target and reuse its tree/normals across deskew rounds.

    Correspondences, normal PCA, Jacobians and robust systems stay on
    device. Only the 6x6 system, count and squared-error sum cross back per
    iteration. Native small solves, stopping policy and pose updates remain
    on CPU; roundoff-sensitive neighbour/planarity decisions use native repairs.
    """

    def __init__(self, target: FloatArray, options: ScanOdometryOptions) -> None:
        _validate_options(options)
        self._host: FloatArray | None = _coordinates(
            _quantize_geometry(target, options.geometry_quantum_m),
            "odometry target",
        )
        if len(self.host) < 3:
            raise ValueError("CUDA odometry targets require at least three points")
        if self.host.nbytes * 8 > 1024**3:
            raise ValueError("CUDA odometry target exceeds the 1 GiB payload budget")
        self._initialize_target(
            _CudaTree(self.host), options, max(1.0, float(np.max(np.abs(self.host)))),
            execution_identity(),
        )

    @property
    def host(self) -> FloatArray:
        """Return native coordinates, downloading a resident target only on demand."""
        if self._host is None:
            with cp.cuda.Device(self.device_id):
                self._host = _coordinates(cp.asnumpy(self.tree.data), "odometry target")
        return self._host

    def _initialize_target(
        self, tree: _TargetTree, options: ScanOdometryOptions, scale: float, identity: str,
    ) -> None:
        """Initialize shared normals, kernels and counters after target validation."""
        self._pending_normals: _PendingNormals | None = None
        self.options = options
        self.device_id = cp.cuda.runtime.getDevice()
        self.tree = tree
        with cp.cuda.Device(self.device_id):
            self._normals = cp.zeros_like(self.tree.data)
            self._planar = cp.zeros(self.point_count, dtype=cp.bool_)
            self.done = cp.zeros(self.point_count, dtype=cp.bool_)
        self.scale = scale
        self.query_repairs = 0
        self.normal_repairs = 0
        self.fused_calls = 0
        self.system_refresh_calls = 0
        self.normal_rows = 0
        self.registration_repairs = 0
        self.motion_repairs = 0
        self.normal_kernel_calls = 0
        self.async_normal_batches = 0
        self._prepared_normal_rows = 0
        self._all_done = False
        self.identity = identity
        with cp.cuda.Device(self.device_id):
            self.system_kernel = cp.RawKernel(_SYSTEM_SOURCE, "system3", options=("--fmad=false",))
            self.fused_kernel = cp.RawKernel(
                _FUSED_SOURCE,
                "project_query_system3",
                options=("--fmad=false",),
            )
            self.covariance_kernel = cp.RawKernel(
                _COVARIANCE_SOURCE,
                "covariance3",
                options=("--fmad=false",),
            )
            self.normal_kernel = cp.RawKernel(
                _NORMAL_SOURCE,
                "normals3",
                options=("--fmad=false",),
            )

    @property
    def point_count(self) -> int:
        """Return the captured target size without requesting host coordinates."""
        return len(self.tree.data)

    def _reserve_queries(self, count: int) -> None:
        """Allow an alternative spatial index to reserve its query workspace."""

    def __del__(self) -> None:
        # A caller can abandon a target while preprocessing/validation fails.
        # Drain the queued writes before any target or pinned buffers are freed.
        pending = getattr(self, "_pending_normals", None)
        if pending is not None:
            try:
                with cp.cuda.Device(self.device_id):
                    pending.ready.synchronize()
            except Exception:
                pass  # CUDA may already be shutting down during finalization.

    def _prime_normals(self, source_rows: int) -> None:
        """Enqueue one bounded dense batch while the caller prepares its source."""
        if (
            self._all_done
            or self.normal_rows
            or self._pending_normals is not None
            or self.point_count > 65536
            or source_rows < self.point_count / 10
        ):
            return
        with cp.cuda.Device(self.device_id):
            # Allocate host ownership before submitting any asynchronous writes.
            pinned = cp.cuda.alloc_pinned_memory(self.point_count)
            native_flags = np.frombuffer(pinned, dtype=np.uint8, count=self.point_count)
            rows = cp.arange(self.point_count, dtype=cp.int64)
            k = min(self.options.normal_neighbours, self.point_count)
            try:
                flags, distances, indexes = self._launch_normals(rows, self.tree.data, k)
                cp.asnumpy(flags, out=native_flags, blocking=False)
                ready = cp.cuda.Event()
                ready.record()
                self._pending_normals = _PendingNormals(
                    rows,
                    flags,
                    distances,
                    indexes,
                    native_flags,
                    pinned,
                    ready,
                    k,
                )
                self.async_normal_batches += 1
            except Exception:
                cp.cuda.get_current_stream().synchronize()
                self.done[rows] = False
                self._all_done = False
                raise

    def _finish_normal_preparation(self) -> None:
        pending = self._pending_normals
        if pending is None:
            return
        with cp.cuda.Device(self.device_id):
            try:
                pending.ready.synchronize()
                self._repair_normals(pending.rows, pending.k, pending.native_flags)
            except Exception:
                self.done[pending.rows] = False
                self._all_done = False
                raise
            else:
                self._prepared_normal_rows += len(pending.rows)
                self._all_done = self._prepared_normal_rows == self.point_count
            finally:
                self._pending_normals = None

    @property
    def normals(self) -> cp.ndarray:
        """Return the complete normal array, materializing it for explicit readers."""
        with cp.cuda.Device(self.device_id):
            self._ensure_all()
        return self._normals

    @property
    def planar(self) -> cp.ndarray:
        """Return all planarity decisions; registration requests only its hits."""
        with cp.cuda.Device(self.device_id):
            self._ensure_all()
        return self._planar

    def _ensure(self, hit: cp.ndarray) -> None:
        self._finish_normal_preparation()
        if self._all_done:
            return
        missing = cp.unique(hit[~self.done[hit]])
        self._compute_normals(missing)

    def _ensure_all(self) -> None:
        self._finish_normal_preparation()
        if self._all_done:
            return
        if self.normal_rows:
            self._ensure(cp.asarray(np.arange(self.point_count, dtype=np.int64)))
        else:
            self._compute_normals(
                cp.asarray(np.arange(self.point_count, dtype=np.int64)),
                complete=True,
            )

    def _compute_normals(self, missing: cp.ndarray, *, complete: bool = False) -> None:
        k = min(self.options.normal_neighbours, self.point_count)
        for start in range(0, len(missing), 65536):
            rows = missing[start : start + 65536]
            queries = (
                self.tree.data[start : start + len(rows)] if complete else self.tree.data[rows]
            )
            flags, _distances, _indexes = self._launch_normals(rows, queries, k)
            try:
                self._repair_normals(rows, k, cp.asnumpy(flags))
            except Exception:
                self.done[rows] = False
                self._all_done = False
                raise
            self._prepared_normal_rows += len(rows)
        self._all_done = self._prepared_normal_rows == self.point_count

    def _launch_normals(
        self,
        rows: cp.ndarray,
        queries: cp.ndarray,
        k: int,
    ) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray]:
        self.normal_rows += len(rows)
        distances, indexes = self.tree.query(queries, k + 1)
        flags = cp.empty(len(rows), dtype=cp.uint8)
        self.normal_kernel(
            ((len(rows) + 127) // 128,),
            (128,),
            (
                self.tree.data,
                indexes,
                distances,
                rows,
                np.int32(len(rows)),
                np.int32(k),
                np.float64(self.options.max_planarity_ratio),
                self._normals,
                self._planar,
                self.done,
                flags,
            ),
        )
        self.normal_kernel_calls += 1
        return flags, distances, indexes

    def _repair_normals(self, rows: cp.ndarray, k: int, flags: NDArray[np.uint8]) -> None:
        repair = np.flatnonzero(flags)
        if len(repair):
            native_rows = cp.asnumpy(rows[cp.asarray(repair)])
            _, native_indices = self.tree.cpu.query(self.host[native_rows], k=k)
            self.query_repairs += len(repair)
            normals, planar = _normals_from_neighbours(
                self.host,
                native_indices,
                k,
                self.options,
            )
            selected = rows[cp.asarray(repair)]
            self._normals[selected] = cp.asarray(normals)
            self._planar[selected] = cp.asarray(planar)
            self.normal_repairs += len(repair)

    def register(self, source: FloatArray, initial: FloatArray) -> ScanRegistration:
        """Capture native coordinates and execute the shared registration loop."""
        host_source = _coordinates(
            _quantize_geometry(source, self.options.geometry_quantum_m),
            "odometry source",
        )
        return self._register_captured(_CapturedSource(host_source), initial)

    def _register_captured(
        self,
        source: _CapturedSource,
        initial: FloatArray,
    ) -> ScanRegistration:
        """Use validated source coordinates with lazy host repair access.

        Native entry points validate before creating this capture. A device
        pipeline must supply finite FP64 coordinates on this target's device.
        """
        if source.points is not None and source.points.device.id != self.device_id:
            raise ValueError("captured source is on a different CUDA device")
        if self.tree.data.nbytes * 8 + source.nbytes * 12 > 1024**3:
            raise ValueError("CUDA odometry inputs exceed the 1 GiB payload budget")
        transform = np.array(initial, dtype=np.float64, copy=True)
        if transform.shape != (4, 4) or not np.isfinite(transform).all():
            raise ValueError("CUDA initial odometry poses require finite 4x4 matrices")
        opts = self.options
        count, iteration, last_count = 0, 0, 0
        converged, rmse = False, math.inf
        last_hessian: FloatArray | None = None
        uncertain_stopping = False
        scale = max(self.scale, source.maximum())
        with cp.cuda.Device(self.device_id):
            self._finish_normal_preparation()
            points = source.device()
            self._reserve_queries(len(points))
            # Dense overlap amortizes one batched normal query/eigensolve.
            # Sparse sources request only their hits to avoid preparing a large,
            # mostly unused target. This changes work scheduling, not geometry.
            if len(points) >= self.point_count / 10:
                self._ensure_all()
            terms = cp.empty((31, len(points)), dtype=cp.float64)
            uncertain = cp.empty(len(points), dtype=cp.float64)
            moved = cp.empty_like(points)
            distances = cp.empty((len(points), 2), dtype=cp.float64)
            indexes = cp.empty((len(points), 2), dtype=cp.int64)
            pending = cp.empty(len(points), dtype=cp.int64)
            for iteration in range(1, opts.max_iterations + 1):  # noqa: B007
                matrix = cp.asarray(transform)
                error = (
                    128
                    * np.finfo(np.float64).eps
                    * scale
                    * max(
                        1.0,
                        float(np.linalg.norm(transform, ord=1)),
                    )
                )
                if len(points):
                    self.fused_kernel(
                        ((len(points) + 127) // 128,),
                        (128,),
                        (
                            points,
                            matrix,
                            self.tree.data,
                            *self.tree.arrays,
                            self._normals,
                            self._planar,
                            self.done,
                            np.int32(len(points)),
                            np.int32(self.point_count),
                            np.float64(opts.max_correspondence_m),
                            np.float64(opts.huber_delta_m),
                            np.float64(error),
                            uncertain,
                            terms,
                            moved,
                            distances,
                            indexes,
                            pending,
                        ),
                    )
                    self.tree.query_calls += 1
                    self.tree.query_rows += len(points)
                    self.fused_calls += 1
                totals = _host(cp.sum(terms, axis=1))
                refresh = False
                if totals[29]:
                    rows = cp.asnumpy(cp.flatnonzero(uncertain))
                    # Recompute after rare native gate/tie repairs. Ordinary
                    # iterations synchronize only for the reduced system.
                    moved = points @ matrix[:3, :3].T + matrix[:3, 3]
                    distances, indexes = self.tree.query(moved, 2)
                    native_moved = source.native_rows(rows) @ transform[:3, :3].T + transform[:3, 3]
                    distance, index = self.tree.cpu.query(
                        native_moved,
                        k=1,
                        distance_upper_bound=opts.max_correspondence_m,
                    )
                    positions = cp.asarray(rows)
                    distances[positions, 0] = cp.asarray(distance)
                    indexes[positions, 0] = cp.asarray(index)
                    moved[positions] = cp.asarray(native_moved)
                    self.query_repairs += len(rows)
                    valid_hits = indexes[:, 0][
                        (distances[:, 0] < opts.max_correspondence_m)
                        & (indexes[:, 0] < self.point_count)
                    ]
                    self._ensure(valid_hits)
                    refresh = True
                elif totals[30]:
                    # The fused query has already captured the exact moved
                    # coordinates and matches. Compute only newly hit normals,
                    # then refresh the system without repeating spatial queries.
                    self._ensure(pending[pending >= 0])
                    refresh = True
                if refresh:
                    self.system_kernel(
                        ((len(points) + 127) // 128,),
                        (128,),
                        (
                            moved,
                            self.tree.data,
                            self._normals,
                            self._planar,
                            distances,
                            indexes,
                            np.int32(len(points)),
                            np.int32(self.point_count),
                            np.float64(opts.max_correspondence_m),
                            np.float64(opts.huber_delta_m),
                            terms,
                        ),
                    )
                    totals = _host(cp.sum(terms[:29], axis=1))
                    self.system_refresh_calls += 1
                count = int(totals[28])
                if count < opts.min_correspondences:
                    break
                upper = np.triu_indices(6)
                hessian = np.zeros((6, 6))
                hessian[upper] = totals[:21]
                hessian[(upper[1], upper[0])] = totals[:21]
                gradient = totals[21:27]
                rmse = math.sqrt(float(totals[27]) / count)
                try:
                    delta = -np.linalg.solve(hessian, gradient)
                except np.linalg.LinAlgError:
                    break
                last_hessian, last_count = hessian, count
                for norm, tolerance in (
                    (_norm3(delta[3:]), opts.convergence_translation_m),
                    (_norm3(delta[:3]), opts.convergence_rotation_rad),
                ):
                    uncertain_stopping |= abs(norm - tolerance) <= 1e-3 * tolerance
                transform = _exp_se3(delta) @ transform
                if (
                    _norm3(delta[3:]) < opts.convergence_translation_m
                    and _norm3(delta[:3]) < opts.convergence_rotation_rad
                ):
                    converged = True
                    break
        near_motion_boundary = self._near_motion_boundary(transform)
        if uncertain_stopping or near_motion_boundary:
            native = register_point_to_plane(
                source.native(),
                self.host,
                initial,
                opts,
                target_tree=self.tree.cpu,
            )
            self.registration_repairs += 1
            self.motion_repairs += int(near_motion_boundary)
            return replace(native, execution_backend=self.identity)
        eigenvalue = (
            0.0 if last_hessian is None else float(np.linalg.eigvalsh(last_hessian)[0]) / last_count
        )
        return ScanRegistration(
            _canonicalize_motion(transform, opts),
            converged and count >= opts.min_correspondences,
            iteration,
            count,
            rmse,
            eigenvalue,
            execution_backend=self.identity,
        )

    def _near_motion_boundary(self, transform: FloatArray) -> bool:
        """Use the native registration when the requested motion grid is ambiguous."""
        if (
            self.options.motion_translation_quantum_m == 0
            and self.options.motion_rotation_quantum_rad == 0
        ):
            return False
        from scipy.spatial.transform import Rotation

        for values, quantum in (
            (transform[:3, 3], self.options.motion_translation_quantum_m),
            (
                Rotation.from_matrix(transform[:3, :3]).as_rotvec(),
                self.options.motion_rotation_quantum_rad,
            ),
        ):
            if quantum and np.any(np.abs(np.remainder(values / quantum, 1) - 0.5) <= 1e-3):
                return True
        return False

    def statistics(self) -> dict[str, int]:
        """Report actual work, including the two kinds of native repairs."""
        return {
            "serialized_tree_captures": int(self.tree.capture_method == "validated_serialized"),
            "public_node_tree_captures": int(self.tree.capture_method == "public_nodes"),
            "gpu_query_calls": self.tree.query_calls,
            "gpu_query_rows": self.tree.query_rows,
            "scalar_query_repairs": self.query_repairs,
            "scalar_normal_repairs": self.normal_repairs,
            "gpu_fused_system_calls": self.fused_calls,
            "gpu_system_refresh_calls": self.system_refresh_calls,
            "gpu_normal_rows": self.normal_rows,
            "gpu_normal_kernel_calls": self.normal_kernel_calls,
            "gpu_async_normal_batches": self.async_normal_batches,
            "scalar_registration_repairs": self.registration_repairs,
            "scalar_motion_repairs": self.motion_repairs,
        }
