# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing of autoware_ptv3 @5eb9286 (research/ptv3/SPEC.md 3) and the index tables of the device graph.

Ported from ``research/ptv3/scripts/preprocess_ref.py`` (bit-exact numpy re-implementation of
``lib/preprocess/preprocess_kernel.cu``; the host tests check equality on every golden frame). Every step that the
CUDA code does in float32 is done in float32 here:

1. features ``(x, y, z, i')``: ``i' = intensity / 255`` for the uint8-intensity formats (XYZIRC, XYZIRCAEDT; the
   default here, as for the T4 / public samples), raw float intensity otherwise (``points2FeaturesKernel`` L143-195);
2. half-open crop ``min <= p < max`` (``cropKernel`` L197-210);
3. grid coordinates ``int32(floorf(p / vs) - floorf(min / vs))`` (``gridCoord`` L297-306);
4. Morton codes, 3 bits per level, ``(major << 2) | (minor << 1) | z``; order "z": major x, "z-trans": major y
   (``serializeCoord`` L267-282);
5. stable sort of the order-0 codes, first (lowest input index) point per voxel = the voxel feature, voxels in
   Morton order, ``inverse_map`` cropped point -> voxel; keep the first ``max_num_voxels`` (256k) voxels
   (``generateFeatures`` L693-934); a frame with fewer than ``min_num_voxels`` (16k) voxels is **rejected** by the
   node (``ptv3_trt.cpp:940-946``: nothing is published) -- recorded in :attr:`Frame.rejected`, the caller decides;
6. pooling metadata of the 4 stride-2 stages: parent key ``code0 >> 3``, CSR ``indptr``, ``cluster`` (child ->
   parent), pooled grid / codes, the z-trans order of each pooled level and the inverses
   (``generateSerializedPoolingMetadata`` L601-691); level-0 orders are the stable argsort of the codes (the
   graph's two ``autoware::Argsort``);
7. sub-manifold neighbour maps of the 5 stages (:mod:`.neighbours`; mode ``alias-first`` by default, PLAN.md D3).

The device-graph tables (:func:`device_tables`) are pure gather indices into tables that carry an explicit zero
**sentinel row** at the end (C19 / P4 rule: PADDED ``ttnn.embedding`` returns the pad row's contents): attention
patch gathers with the ONNX clip padding, the pooling child lists (segment max as gather + max), the unpool
clusters, the per-stage sparse-conv im2col indices (C13 ``im2col_index``) and the BEV scatter-max candidate lists of
the det head. Every table is a function of the frame alone, so the CPU reference and the device consume identical
tables (the PCC protocol of SPEC 6).

numpy only; no side effects on import.
"""
from __future__ import annotations

import hashlib
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence

import numpy as np

from ..reference.config import (BEV_CELL_CANDIDATES, BEV_SIZE, NUM_STAGES, PATCH_SIZE, POOL_CHILDREN, SPATIAL_SHAPES,
                                STAGE_CAPACITIES, PreprocessConfig)
from .neighbours import neighbour_map, to_pairs

__all__ = ["points_to_features", "serialize", "grid_coord", "Voxels", "voxelize", "PoolStage", "pooling_metadata",
           "Level", "Frame", "prepare_frame", "attention_index", "pool_index", "unpool_index", "bev_index",
           "im2col_index", "DeviceTables", "device_tables", "table_digests", "sha256_array", "padded_count"]


# ---------------------------------------------------------------------------------------------- steps 1-5
def points_to_features(points: np.ndarray, intensity_is_uint8: bool = True) -> np.ndarray:
    """``points2FeaturesKernel``: ``(N, >= 4)`` -> float32 ``(N, 4)`` = x, y, z, i'."""
    pts = np.asarray(points)
    if pts.ndim != 2 or pts.shape[1] < 4:
        raise ValueError(f"points must be (N, >= 4) x, y, z, intensity, got {pts.shape}")
    feats = np.empty((pts.shape[0], 4), np.float32)
    feats[:, :3] = pts[:, :3].astype(np.float32)
    if intensity_is_uint8:
        feats[:, 3] = pts[:, 3].astype(np.uint8).astype(np.float32) / np.float32(255.0)
    else:
        feats[:, 3] = pts[:, 3].astype(np.float32)
    return feats


def _spread_bits(v: np.ndarray, depth: int) -> np.ndarray:
    v = np.asarray(v).astype(np.int64)
    out = np.zeros_like(v)
    for i in range(depth):
        out |= ((v >> i) & 1) << (3 * i)
    return out


def serialize(gx: np.ndarray, gy: np.ndarray, gz: np.ndarray, depth: int, transposed: bool) -> np.ndarray:
    """``serializeCoord``: int64 Morton code (major << 2 | minor << 1 | z per bit level)."""
    major, minor = (gy, gx) if transposed else (gx, gy)
    return (_spread_bits(major, depth) << 2) | (_spread_bits(minor, depth) << 1) | _spread_bits(gz, depth)


def grid_coord(xyz: np.ndarray, cfg: PreprocessConfig) -> np.ndarray:
    """``gridCoord``: int32 ``floor(p / vs) - coord_min`` in float32."""
    xyz = np.asarray(xyz, np.float32)
    vs = [np.float32(v) for v in cfg.voxel_size]
    cmin = cfg.coord_min
    out = np.empty((xyz.shape[0], 3), np.int32)
    for i in range(3):
        out[:, i] = (np.floor(xyz[:, i] / vs[i]) - np.float32(cmin[i])).astype(np.int32)
    return out


@dataclass
class Voxels:
    feat: np.ndarray                 # float32 (N0, 4): the first in-range point of each voxel
    grid_coord: np.ndarray           # int32 (N0, 3)
    serialized_code: np.ndarray      # int64 (2, N0): z / z-trans Morton codes
    num_unique_voxels: int           # before the max_num_voxels clip
    crop_mask: np.ndarray            # bool (N_in,)
    cropped_index: np.ndarray        # int64 (N_crop,): input row of each cropped point
    inverse_map: np.ndarray          # int64 (N_crop,): voxel id of each cropped point (may be >= N0 when clipped)
    representative: np.ndarray       # int64 (N0,): cropped index of each voxel's point

    @property
    def num_voxels(self) -> int:
        return int(self.feat.shape[0])


def _ops(impl: str):
    """``(serialize, argsort_stable, inverse_permutation)`` of ``impl``: ``"numpy"`` (the oracles) or ``"numba"``
    (the bit-exact twins of :mod:`.pre_nb`, OPT round 4; numpy when numba is missing)."""
    if impl == "numba":
        from . import pre_nb

        if pre_nb.available():
            return pre_nb.serialize, pre_nb.argsort_stable, pre_nb.inverse_permutation
    elif impl != "numpy":
        raise ValueError(f"pre impl must be 'numpy' or 'numba', got {impl!r}")
    return (serialize, lambda k: np.argsort(k, kind="stable").astype(np.int64),
            lambda p: np.argsort(p).astype(np.int64))


def voxelize(feats: np.ndarray, cfg: PreprocessConfig, impl: str = "numpy") -> Voxels:
    """``generateFeatures`` (steps 2-5 of the module docstring). ``impl``: see :func:`_ops`."""
    ser, asort, _ = _ops(impl)
    f32 = np.float32
    r = [f32(v) for v in cfg.point_cloud_range]
    x, y, z = feats[:, 0], feats[:, 1], feats[:, 2]
    mask = (x >= r[0]) & (x < r[3]) & (y >= r[1]) & (y < r[4]) & (z >= r[2]) & (z < r[5])
    cropped_idx = np.nonzero(mask)[0].astype(np.int64)
    cropped = feats[cropped_idx]
    depth = cfg.serialization_depth
    g = grid_coord(cropped[:, :3], cfg)
    code0 = ser(g[:, 0], g[:, 1], g[:, 2], depth, False)
    order = asort(code0)                                       # cub radix SortPairs is stable
    sorted_codes = code0[order]
    run_start = np.ones(sorted_codes.shape[0], bool)
    run_start[1:] = sorted_codes[1:] != sorted_codes[:-1]
    uid = np.cumsum(run_start) - 1
    num_unique = int(uid[-1]) + 1 if uid.size else 0
    rep = order[np.nonzero(run_start)[0]]                      # lowest cropped index of each voxel
    keep = min(num_unique, cfg.max_num_voxels)
    voxel_feat = cropped[rep[:keep]].copy()
    inverse_map = np.empty(cropped.shape[0], np.int64)
    inverse_map[order] = uid
    gv = grid_coord(voxel_feat[:, :3], cfg)                    # recomputed from the representative point
    codes = np.stack([ser(gv[:, 0], gv[:, 1], gv[:, 2], depth, False),
                      ser(gv[:, 0], gv[:, 1], gv[:, 2], depth, True)]).astype(np.int64)
    return Voxels(voxel_feat.astype(np.float32), gv.astype(np.int32), codes, num_unique, mask, cropped_idx,
                  inverse_map, rep[:keep].astype(np.int64))


# ---------------------------------------------------------------------------------------------- step 6
@dataclass
class PoolStage:
    indices: np.ndarray              # int64 (N_s,): identity (children are contiguous)
    indptr: np.ndarray               # int64 (N_{s+1} + 1,)
    head_indices: np.ndarray         # int64 (N_{s+1},)
    cluster: np.ndarray              # int64 (N_s,): child -> parent
    grid_coord: np.ndarray           # int32 (N_{s+1}, 3)
    serialized_order: np.ndarray     # int64 (2, N_{s+1})
    serialized_inverse: np.ndarray   # int64 (2, N_{s+1})
    serialized_code: np.ndarray      # int64 (2, N_{s+1})

    @property
    def in_count(self) -> int:
        return int(self.cluster.shape[0])

    @property
    def out_count(self) -> int:
        return int(self.head_indices.shape[0])


def pooling_metadata(grid: np.ndarray, codes: np.ndarray, strides: Sequence[int], num_orders: int = 2,
                     impl: str = "numpy") -> List[PoolStage]:
    """``generateSerializedPoolingMetadata``; the input level must be strictly ascending in the order-0 code."""
    _, asort, ainv = _ops(impl)
    n = grid.shape[0]
    if n and not np.all(np.diff(codes[0]) > 0):
        raise ValueError("the input level must be strictly ascending in the order-0 code")
    cur_order = np.stack([np.arange(n, dtype=np.int64)]
                         + [asort(codes[o]) for o in range(1, num_orders)])
    cur_grid, cur_codes = grid, codes
    stages: List[PoolStage] = []
    for stride in strides:
        depth = int(stride - 1).bit_length()
        key = cur_codes[0] >> (3 * depth)
        flag = np.ones(key.shape[0], bool)
        flag[1:] = key[1:] != key[:-1]
        cluster = (np.cumsum(flag) - 1).astype(np.int64)
        m = int(cluster[-1]) + 1 if cluster.size else 0
        heads = np.nonzero(flag)[0].astype(np.int64)
        indptr = np.concatenate([heads, np.array([key.shape[0]], np.int64)])
        g_out = (cur_grid[heads] >> depth).astype(np.int32)
        c_out = (cur_codes[:, heads] >> (3 * depth)).astype(np.int64)
        order_out = np.empty((num_orders, m), np.int64)
        inv_out = np.empty((num_orders, m), np.int64)
        order_out[0] = np.arange(m)
        inv_out[0] = np.arange(m)
        for o in range(1, num_orders):            # walk the level in z-trans rank, compact parent-run heads
            seg_walk = cluster[cur_order[o]]
            f = np.ones(seg_walk.shape[0], bool)
            f[1:] = seg_walk[1:] != seg_walk[:-1]
            order_out[o] = seg_walk[f]
            inv_out[o] = ainv(order_out[o])
        stages.append(PoolStage(np.arange(key.shape[0], dtype=np.int64), indptr, heads, cluster, g_out, order_out,
                                inv_out, c_out))
        cur_grid, cur_codes, cur_order = g_out, c_out, order_out
    return stages


# ---------------------------------------------------------------------------------------------- frame
@dataclass
class Level:
    """One encoder stage: voxel coordinates, the two serialization orders and inverses, the neighbour map."""
    grid: np.ndarray                 # int32 (N_s, 3)
    order: np.ndarray                # int64 (2, N_s)
    inverse: np.ndarray              # int64 (2, N_s)
    nmap: np.ndarray                 # int32 (N_s, 27)

    @property
    def count(self) -> int:
        return int(self.grid.shape[0])


@dataclass
class Frame:
    """Everything the network consumes for one frame, plus what the post-processing needs."""
    num_points: int
    voxels: Voxels
    pools: List[PoolStage]
    levels: List[Level]
    neighbour_mode: str
    rejected: bool                   # node @5eb9286 publishes nothing (< min_num_voxels)
    clipped: bool                    # more than max_num_voxels unique voxels (the node clips, so do we)
    overflow: bool = False           # a stage count exceeds the device capacity (STAGE_CAPACITIES)
    timing_ms: Dict[str, float] = field(default_factory=dict)

    @property
    def counts(self) -> List[int]:
        return [lv.count for lv in self.levels]

    @property
    def clusters(self) -> List[np.ndarray]:
        return [p.cluster for p in self.pools]

    def encoder_inputs(self) -> Dict[str, np.ndarray]:
        """The 27 inputs of ``ptv3_encoder.onnx`` (names and dtypes as in the graph)."""
        v = self.voxels
        enc: Dict[str, np.ndarray] = {"grid_coord": v.grid_coord, "feat": v.feat, "serialized_code": v.serialized_code}
        for s, st in enumerate(self.pools):
            for k in ("indices", "indptr", "head_indices", "grid_coord", "serialized_order", "serialized_inverse"):
                a = getattr(st, k)
                enc[f"serialized_pooling_{s}_{k}"] = a.astype(np.int32 if k == "grid_coord" else np.int64)
        return enc

    def seg_head_inputs(self, point_feats: Sequence[np.ndarray]) -> Dict[str, np.ndarray]:
        """The 15 inputs of ``ptv3_seg3d_head.onnx``."""
        d = {f"point_feat_{s}": np.asarray(point_feats[s], np.float32) for s in range(NUM_STAGES)}
        for s, st in enumerate(self.pools):
            d[f"pooling_cluster_{s}"] = st.cluster
        for s in (1, 2):
            st = self.pools[s]
            d[f"serialized_pooling_{s}_serialized_order"] = st.serialized_order
            d[f"serialized_pooling_{s}_serialized_inverse"] = st.serialized_inverse
            d[f"serialized_pooling_{s}_grid_coord"] = st.grid_coord.astype(np.int32)
        return d

    def det_head_inputs(self, point_feats: Sequence[np.ndarray]) -> Dict[str, np.ndarray]:
        """The 4 inputs of ``ptv3_det3d_head.onnx`` (``point_grid_coord_3`` = stage-3 grid, ptv3_trt.cpp:678-683)."""
        return {"point_feat_3": np.asarray(point_feats[3], np.float32),
                "point_feat_4": np.asarray(point_feats[4], np.float32),
                "pooling_cluster_3": self.pools[3].cluster, "point_grid_coord_3": self.pools[2].grid_coord}

    def describe(self) -> Dict[str, Any]:
        return {"num_points": self.num_points, "points_in_range": int(self.voxels.crop_mask.sum()),
                "num_unique_voxels": self.voxels.num_unique_voxels, "stage_counts": self.counts,
                "neighbour_mode": self.neighbour_mode, "rejected": self.rejected, "clipped": self.clipped,
                "overflow": self.overflow}


_POOLS: Dict[int, Any] = {}
_POOLS_LOCK = threading.Lock()


def _pool(n: int):
    """A process-wide ``ThreadPoolExecutor`` of ``n`` workers (created once per size)."""
    with _POOLS_LOCK:
        p = _POOLS.get(n)
        if p is None:
            from concurrent.futures import ThreadPoolExecutor

            p = _POOLS[n] = ThreadPoolExecutor(max_workers=n, thread_name_prefix="ptv3-maps")
        return p


def prepare_frame(points: np.ndarray, cfg: Optional[PreprocessConfig] = None, *, neighbour_mode: str = "alias-first",
                  intensity_is_uint8: bool = True, capacities: Optional[Sequence[int]] = STAGE_CAPACITIES,
                  neighbour_impl: str = "numpy", pre_impl: str = "numpy", threads: int = 1) -> Frame:
    """Steps 1-7 for one cloud ``(N, >= 4)`` x, y, z, intensity (base_link). ``neighbour_impl``: ``"numpy"`` (the
    oracles of :mod:`.neighbours`) or ``"numba"`` (the bit-exact twins of :mod:`.neighbours_nb`; the served path,
    ``PTV3_HOST_MAPS``; falls back to numpy when numba is missing). ``pre_impl``: the serialization / sorts of steps
    4-6 as ``"numpy"`` or the bit-exact ``"numba"`` twins of :mod:`.pre_nb` (``PTV3_HOST_PRE``). ``threads`` > 1:
    the 5 stage neighbour maps (independent; the numba twins release the GIL) run concurrently on a shared thread
    pool of that size (``PTV3_HOST_THREADS``; the same maps)."""
    if neighbour_impl not in ("numpy", "numba"):
        raise ValueError(f"neighbour_impl must be 'numpy' or 'numba', got {neighbour_impl!r}")
    nmap_fn = neighbour_map
    if neighbour_impl == "numba":
        from . import neighbours_nb

        if neighbours_nb.available():
            nmap_fn = neighbours_nb.neighbour_map_nb
    cfg = cfg or PreprocessConfig()
    t0 = time.perf_counter()
    feats = points_to_features(points, intensity_is_uint8)
    _, asort, ainv = _ops(pre_impl)
    vox = voxelize(feats, cfg, pre_impl)
    t1 = time.perf_counter()
    pools = pooling_metadata(vox.grid_coord, vox.serialized_code, cfg.pooling_strides, cfg.num_orders, pre_impl)
    order0 = np.stack([asort(vox.serialized_code[o]) for o in range(cfg.num_orders)]).astype(np.int64)
    inv0 = np.stack([ainv(order0[o]) for o in range(cfg.num_orders)]).astype(np.int64)
    grids = [vox.grid_coord] + [p.grid_coord for p in pools]
    orders = [(order0, inv0)] + [(p.serialized_order, p.serialized_inverse) for p in pools]
    t2 = time.perf_counter()
    if int(threads) > 1:
        nmaps = list(_pool(int(threads)).map(lambda s: nmap_fn(grids[s], SPATIAL_SHAPES[s], neighbour_mode),
                                             range(NUM_STAGES)))
    else:
        nmaps = [nmap_fn(grids[s], SPATIAL_SHAPES[s], neighbour_mode) for s in range(NUM_STAGES)]
    levels = [Level(grids[s].astype(np.int32), orders[s][0], orders[s][1], nmaps[s]) for s in range(NUM_STAGES)]
    t3 = time.perf_counter()
    counts = [lv.count for lv in levels]
    overflow = bool(capacities is not None and any(c > cap for c, cap in zip(counts, capacities)))
    return Frame(int(np.asarray(points).shape[0]), vox, pools, levels, neighbour_mode,
                 rejected=vox.num_unique_voxels < cfg.min_num_voxels, clipped=vox.num_unique_voxels > cfg.max_num_voxels,
                 overflow=overflow,
                 timing_ms={"voxelize": (t1 - t0) * 1e3, "pooling": (t2 - t1) * 1e3, "neighbours": (t3 - t2) * 1e3})


# ---------------------------------------------------------------------------------------------- device tables
def padded_count(n: int) -> int:
    """Rows of the serialized attention of ``n`` voxels: ``ceil(n / 512) * 512``."""
    return -(-int(n) // PATCH_SIZE) * PATCH_SIZE


def attention_index(order: np.ndarray, n: int, capacity: int) -> np.ndarray:
    """uint32 ``(capacity,)``: row ``j`` of the patch-major sequence reads voxel ``order[min(j, n - 1)]`` (the ONNX
    clip padding: pad tokens are copies of the last voxel in curve order, no mask). Rows past the padded count
    (whole extra patches, which nothing reads back) repeat that voxel too."""
    if n <= 0:
        return np.zeros(capacity, np.uint32)
    j = np.minimum(np.arange(capacity), n - 1)
    return np.asarray(order, np.int64)[j].astype(np.uint32)


def pool_index(indptr: np.ndarray, n_in: int, capacity_out: int, *, sentinel: int,
               children: int = POOL_CHILDREN) -> np.ndarray:
    """uint32 ``(capacity_out, children)``: the child rows of each parent (``SegmentCSR`` max as gather + max);
    short lists repeat their first child (max is idempotent), rows past the parent count read ``sentinel``."""
    indptr = np.asarray(indptr, np.int64)
    m = indptr.shape[0] - 1
    lens = indptr[1:] - indptr[:-1]
    if m and (lens.min() < 1 or lens.max() > children):
        raise ValueError(f"pool segments must hold 1..{children} children, got {lens.min()}..{lens.max()}")
    out = np.full((capacity_out, children), sentinel, np.uint32)
    if m:
        k = np.arange(children)[None, :]
        idx = indptr[:-1, None] + np.where(k < lens[:, None], k, 0)
        out[:m] = idx.astype(np.uint32)
    if m and int(indptr[-1]) != n_in:
        raise ValueError("indptr does not cover the input level")
    return out


def unpool_index(cluster: np.ndarray, capacity: int, *, sentinel: int) -> np.ndarray:
    """uint32 ``(capacity,)``: child row -> parent row (the seg / det heads' ``Gather(cluster)``); padding rows read
    ``sentinel``."""
    out = np.full(capacity, sentinel, np.uint32)
    out[:len(cluster)] = np.asarray(cluster, np.int64).astype(np.uint32)
    return out


def bev_index(grid3: np.ndarray, *, sentinel: int, size: int = BEV_SIZE,
              candidates: int = BEV_CELL_CANDIDATES) -> np.ndarray:
    """uint32 ``(size * size, candidates)``: the stage-3 voxel rows of each BEV cell ``y * size + x`` (cells with
    ``0 <= x, y < size``, the graph's mask), padded with ``sentinel`` (a zero row). The det head's
    ``ScatterElements(reduction=max)`` into a zero canvas of non-negative (ReLU) features is then ``max`` over the
    candidates (the zero row stands for the canvas' initial zeros)."""
    g = np.asarray(grid3, np.int64)
    out = np.full((size * size, candidates), sentinel, np.uint32)
    ok = (g[:, 0] >= 0) & (g[:, 0] < size) & (g[:, 1] >= 0) & (g[:, 1] < size)
    rows = np.nonzero(ok)[0]
    cell = g[rows, 1] * size + g[rows, 0]
    order = np.argsort(cell, kind="stable")
    cell, rows = cell[order], rows[order]
    start = np.ones(len(cell), bool)
    start[1:] = cell[1:] != cell[:-1]
    first = np.maximum.accumulate(np.where(start, np.arange(len(cell)), 0))
    slot = np.arange(len(cell)) - first
    if len(slot) and slot.max() >= candidates:
        raise ValueError(f"a BEV cell holds {slot.max() + 1} stage-3 voxels (> {candidates})")
    out[cell, slot] = rows.astype(np.uint32)
    return out


def im2col_index(nmap: np.ndarray, capacity: int, *, sentinel: int) -> np.ndarray:
    """C13 tile-ordered im2col index ``(capacity / 32, 27 * 32)`` of a stage's neighbour map (C28 input)."""
    from ..ttaw.sparse import im2col_index as _im2col
    return _im2col(nmap, sentinel=sentinel, rows=capacity)


@dataclass
class DeviceTables:
    """The per-frame RT-dev index tables of the device graph (uint32; sentinel = the capacity row of each table)."""
    capacities: List[int]
    counts: List[int]
    attention: List[np.ndarray]      # per stage: (2, cap_s) patch gathers (orders z, z-trans)
    inverse: List[np.ndarray]        # per stage: (2, cap_s) original row -> sequence row (rows >= N_s read row 0)
    neighbours: List[np.ndarray]     # per stage: (cap_s, 27) gather rows, sentinel cap_s
    pool: List[np.ndarray]           # stages 1..4: (cap_s, 8) child rows into stage s-1, sentinel cap_{s-1}
    unpool: List[np.ndarray]         # stages 0..3: (cap_s,) parent rows into stage s+1, sentinel cap_{s+1}
    bev: Optional[np.ndarray] = None  # (65536, 9) stage-3 rows, sentinel cap_3

    def nbytes(self) -> int:
        arrs = self.attention + self.inverse + self.neighbours + self.pool + self.unpool
        return int(sum(a.nbytes for a in arrs) + (self.bev.nbytes if self.bev is not None else 0))


def device_tables(frame: Frame, capacities: Sequence[int] = STAGE_CAPACITIES, *, with_bev: bool = True) -> DeviceTables:
    """Every gather table of the device graph for ``frame`` (fixed shapes from ``capacities``)."""
    caps = [int(c) for c in capacities]
    counts = frame.counts
    if any(n > c for n, c in zip(counts, caps)):
        raise ValueError(f"stage counts {counts} exceed the capacities {caps}")
    att, inv, nb, pool, unpool = [], [], [], [], []
    for s, lv in enumerate(frame.levels):
        n, cap = counts[s], caps[s]
        att.append(np.stack([attention_index(lv.order[o], n, cap) for o in range(lv.order.shape[0])]))
        iv = np.zeros((lv.inverse.shape[0], cap), np.uint32)
        iv[:, :n] = lv.inverse.astype(np.uint32)
        inv.append(iv)
        g = np.full((cap, lv.nmap.shape[1]), cap, np.uint32)
        g[:n] = np.where(lv.nmap >= 0, lv.nmap, cap).astype(np.uint32)
        nb.append(g)
    for s, st in enumerate(frame.pools):
        pool.append(pool_index(st.indptr, counts[s], caps[s + 1], sentinel=caps[s]))
        unpool.append(unpool_index(st.cluster, caps[s], sentinel=caps[s + 1]))
    bev = bev_index(frame.pools[2].grid_coord, sentinel=caps[3]) if with_bev else None
    return DeviceTables(caps, counts, att, inv, nb, pool, unpool, bev)


# ---------------------------------------------------------------------------------------------- digests
def sha256_array(a: np.ndarray) -> str:
    a = np.ascontiguousarray(a)
    return hashlib.sha256(a.view(np.uint8).reshape(-1).tobytes()).hexdigest()


def table_digests(frame: Frame) -> Dict[str, str]:
    """sha256 of every host table in the layout of the research goldens' ``host_table_sha256``
    (research/ptv3/public_data/scripts/run_reference.py ``table_digests``): a TT run proves it consumed the tables
    the golden was made with."""
    enc = frame.encoder_inputs()
    d = {f"enc_in.{k}": sha256_array(v) for k, v in sorted(enc.items())}
    d.update({f"pooling_cluster_{s}": sha256_array(p.cluster) for s, p in enumerate(frame.pools)})
    d.update({f"subm_pair.stage{s}": sha256_array(to_pairs(lv.nmap)) for s, lv in enumerate(frame.levels)})
    return d