vt1200/harness/udpstream/tilecost.py

"""Build a 32x32 1-bpp tile dictionary and quantise frames into tile indices.

Scheme C replaces row strips with a fixed dictionary of ``k`` 32x32 tiles plus
one index byte per tile, so a full screen is 1,280 index bytes on the wire
instead of 20,480 bytes of rows.  This module is the offline half: it builds the
dictionary from the clip's frames and turns a frame into indices.  The
terminal's half is one ``PIXBLT B,XY`` per tile from ``dict + index*128`` with
``SPTCH=32`` and ``DYDX=(32,32)``; see the design's §8.3.

Geometry
--------

The screen is 1280x1024, one bit per pixel.  A source row is therefore 160 bytes
(:data:`harness.udpstream.vidsrc.ROW_BYTES`) and a 32x32 tile is 32 rows of
**four** bytes -- 32 px is two 16-bit words -- for 128 bytes: row-major, two
bytes per word, **LSB first** within each byte (bit 0 is the leftmost pixel,
the scanout order), the same bit order
:meth:`harness.udpstream.vidsrc.Rows.pack_row` produces.  The grid is 40 columns
by 32 rows: 1,280 tiles.

A frame buffer is whole 160-byte rows and whole 32-row tile bands.  A buffer
shorter than the 1024-row screen is the **top** of the screen: :func:`quantise`
still returns the full 40x32 index frame the terminal's ``TILES_FULL`` expects,
and the rows the buffer does not cover are quantised from the all-zero (black)
tile.  That is the only reading that satisfies both the wire contract (a full
frame is always 1,280 indices) and the plan's 128-row test frame, which covers
160 of the 1,280 tiles; it is pinned by a test rather than left implicit.

The dictionary
--------------

:func:`build_dictionary` runs Lloyd's algorithm over every tile of every
supplied frame with Hamming distance as the metric -- the right metric here,
because a displaced wire bit is one wrong pixel.  It is deterministic end to
end:

* only **distinct** tiles are candidates (repeats add nothing to a vote), and
  candidates are named by content, so the result does not depend on the order
  the frames are supplied in;
* the initial centroids are the ``k`` candidates whose ``blake2b(seed, tile)``
  digest sorts lowest.  That is a seeded choice, but unlike
  :meth:`random.Random.sample` it is stable across Python versions, and the
  dictionary is meant to be committed once and reproduced;
* a tile joins the nearest centroid, ties to the **lowest** index;
* a centroid is the per-pixel majority of its members, which is the Hamming
  centroid that minimises total distance.  A pixel with exactly half its members
  set stays 0: ties are black;
* an empty cluster keeps its previous centroid;
* iteration stops when the centroids stop moving, and in any case after
  :data:`ITERATIONS` (8) passes;
* the centroids are finally sorted as byte strings, so the dictionary's *order*
  is canonical as well as its contents.

With fewer distinct candidate tiles than ``k`` the remaining entries are the
all-zero tile: unused slots are black, never invented.  A short frame is not
padded here -- padding the training set with the 896 unseen black tile rows
would outvote the picture -- so build the dictionary from frames that cover
what the clip actually shows (the full-screen frames it will send).

Cost: one quantise is ``tiles x k`` Hamming distances (327,680 for a full screen
at ``k=256``), but repeated tiles share one answer, which is what makes a frame
of flat panels cheap.  The dictionary is built offline once per clip.
"""

import hashlib

from harness.udpstream.vidsrc import ROW_BYTES, SCREEN_ROWS

__all__ = [
    "ITERATIONS", "MAX_TILES", "SCREEN_TILE_ROWS", "TILES_PER_SCREEN",
    "TILE_BYTES", "TILE_COLS", "TILE_PIXELS", "TILE_ROW_BYTES",
    "build_dictionary", "quantise",
]

#: Side of a tile, in pixels (and rows).
TILE_PIXELS = 32

#: Bytes per tile row: 32 px packed one bit per pixel.
TILE_ROW_BYTES = TILE_PIXELS // 8                       # 4

#: A whole tile: 32 rows of four bytes, row-major.
TILE_BYTES = TILE_PIXELS * TILE_ROW_BYTES               # 128

#: Tiles across one 160-byte source row.
TILE_COLS = ROW_BYTES // TILE_ROW_BYTES                 # 40

#: Tile bands down the screen.
SCREEN_TILE_ROWS = SCREEN_ROWS // TILE_PIXELS           # 32

#: Indices in one full-screen ``TILES_FULL`` frame.
TILES_PER_SCREEN = TILE_COLS * SCREEN_TILE_ROWS         # 1,280

#: Most dictionary entries: an index travels as one byte.
MAX_TILES = 256

#: Most Lloyd passes, whether or not the centroids have settled.
ITERATIONS = 8

_TILE_BITS = TILE_BYTES * 8
_ALL = (1 << _TILE_BITS) - 1


def _tile_rows(frame):
    """Return how many 32-row tile bands ``frame`` holds, or raise ValueError.

    The buffer must be whole 160-byte rows and whole 32-row bands, and at most
    the 1024-row screen: a partial row or band has no tile to name, and rows
    below the screen would address memory the terminal does not have.
    """
    try:
        length = len(frame)
    except TypeError:
        raise TypeError("a frame is bytes-like, got %s"
                        % type(frame).__name__)
    if length % ROW_BYTES:
        raise ValueError("a frame is whole %d-byte rows, got %d bytes"
                         % (ROW_BYTES, length))
    rows = length // ROW_BYTES
    if rows % TILE_PIXELS:
        raise ValueError("a frame is whole %d-row tiles, got %d rows"
                         % (TILE_PIXELS, rows))
    if not 1 <= rows <= SCREEN_ROWS:
        raise ValueError("a frame is 1..%d rows of %d bytes, got %d"
                         % (SCREEN_ROWS, ROW_BYTES, rows))
    return rows // TILE_PIXELS


def _frame_tiles(frame):
    """Return ``frame``'s tiles as ints, row-major over the tile grid.

    Tile (tx, ty) is rows ``ty*32``..``+31``, bytes ``tx*4``..``tx*4+4`` of each
    row, concatenated in that order -- exactly the 128 bytes the blitter reads
    from ``dict + index*128``.  An int is the cheapest Hamming carrier:
    ``(a ^ b).bit_count()`` is one C call.
    """
    tiles = []
    for ty in range(_tile_rows(frame)):
        top = ty * TILE_PIXELS * ROW_BYTES
        for tx in range(TILE_COLS):
            left = tx * TILE_ROW_BYTES
            tiles.append(int.from_bytes(
                b"".join(frame[top + row * ROW_BYTES + left:
                               top + row * ROW_BYTES + left + TILE_ROW_BYTES]
                         for row in range(TILE_PIXELS)), "big"))
    return tiles


def _dictionary_entries(dictionary):
    """Return the dictionary as ints, or raise ValueError on a bad length."""
    try:
        length = len(dictionary)
    except TypeError:
        raise TypeError("a dictionary is bytes-like, got %s"
                        % type(dictionary).__name__)
    if length % TILE_BYTES:
        raise ValueError("a dictionary is whole %d-byte tiles, got %d bytes"
                         % (TILE_BYTES, length))
    count = length // TILE_BYTES
    if not 1 <= count <= MAX_TILES:
        raise ValueError("a dictionary is 1..%d tiles of %d bytes (an index is "
                         "one byte), got %d"
                         % (MAX_TILES, TILE_BYTES, count))
    return [int.from_bytes(dictionary[i * TILE_BYTES:(i + 1) * TILE_BYTES],
                           "big")
            for i in range(count)]


def _nearest(tile, entries):
    """Index of the entry with the smallest Hamming distance; ties lower."""
    best = 0
    best_distance = (tile ^ entries[0]).bit_count()
    if best_distance == 0:
        return 0
    for index in range(1, len(entries)):
        distance = (tile ^ entries[index]).bit_count()
        if distance == 0:
            return index        # nothing can beat zero, and later ties lose
        if distance < best_distance:
            best, best_distance = index, distance
    return best


def _majority(tiles):
    """Per-bit majority of ``tiles`` as one int; a tie (half set) stays 0.

    Counts are accumulated bit-sliced: plane *j* holds bit *j* of the count at
    every pixel position, so one addition is ``~log2(n)`` big-int operations and
    the threshold is a single borrow-propagating subtraction.  Counting pixels
    one at a time would be 1,024 Python operations per member.
    """
    planes = []
    for tile in tiles:
        carry = tile
        shift = 0
        while carry:
            if shift == len(planes):
                planes.append(0)
            carry, planes[shift] = planes[shift] & carry, planes[shift] ^ carry
            shift += 1

    need = len(tiles) // 2 + 1          # strict majority
    borrow = 0                          # borrow out of count - need
    for shift in range(max(len(planes), need.bit_length())):
        plane = planes[shift] if shift < len(planes) else 0
        subtrahend = _ALL if (need >> shift) & 1 else 0
        borrow = ((~plane & _ALL) & (subtrahend | borrow)) | (subtrahend
                                                              & borrow)
    # A pixel whose subtraction ended without a borrow counts at least ``need``.
    return ~borrow & _ALL


def _initial_centroids(candidates, k, seed):
    """Choose ``k`` distinct candidates as a deterministic function of content.

    With at most ``k`` candidates every one is kept (the rest of the dictionary
    is padded black).  Otherwise the choice is the ``k`` lowest
    ``blake2b(seed, tile)`` digests.  Seeded, but stable across Python versions
    and independent of frame order, which matters because the dictionary is
    committed output.  ``random.Random.sample`` guarantees neither.
    """
    if len(candidates) <= k:
        return candidates + [0] * (k - len(candidates))
    prefix = str(seed).encode("ascii")

    def key(tile):
        digest = hashlib.blake2b(prefix)
        digest.update(tile.to_bytes(TILE_BYTES, "big"))
        return digest.digest()

    return sorted(candidates, key=key)[:k]


def build_dictionary(frames, k=256, seed=0):
    """Return ``k`` tiles of :data:`TILE_BYTES` bytes, deterministically ordered.

    ``frames`` are packed 1-bpp buffers as :func:`quantise` takes them; every
    tile of every frame feeds the k-means.  ``k`` defaults to 256, the most an
    index byte can name.  ``seed`` chooses the initial centroids; the same
    ``seed``, frames and ``k`` always give byte-identical output.
    """
    if not isinstance(k, int) or isinstance(k, bool) or not 1 <= k <= MAX_TILES:
        raise ValueError("k is 1..%d (an index is one byte), got %r"
                         % (MAX_TILES, k))
    frames = list(frames)
    if not frames:
        raise ValueError("build_dictionary needs at least one frame")

    tiles = []
    for frame in frames:
        tiles.extend(_frame_tiles(frame))

    seen = set()
    candidates = []
    for tile in tiles:
        if tile not in seen:
            seen.add(tile)
            candidates.append(tile)

    centroids = _initial_centroids(candidates, k, seed)
    for _ in range(ITERATIONS):
        buckets = [[] for _ in range(k)]
        nearest = {}
        for tile in tiles:
            index = nearest.get(tile)
            if index is None:
                index = _nearest(tile, centroids)
                nearest[tile] = index
            buckets[index].append(tile)
        updated = [centroids[i] if not bucket else _majority(bucket)
                   for i, bucket in enumerate(buckets)]
        if updated == centroids:
            break
        centroids = updated

    centroids.sort()
    return b"".join(tile.to_bytes(TILE_BYTES, "big") for tile in centroids)


def quantise(frame, dictionary):
    """Return the frame's :data:`TILES_PER_SCREEN` indices, row-major.

    Each index is the dictionary entry with the smallest Hamming distance to
    that 32x32 tile, ties to the lowest index, so the terminal can paint the
    whole screen from the returned bytes alone.  A frame shorter than 1024 rows
    is the top of the screen; the rest quantises the all-zero tile.
    """
    entries = _dictionary_entries(dictionary)
    tiles = _frame_tiles(frame)

    nearest = {}
    indices = bytearray(TILES_PER_SCREEN)
    for position, tile in enumerate(tiles):
        index = nearest.get(tile)
        if index is None:
            index = _nearest(tile, entries)
            nearest[tile] = index
        indices[position] = index

    black = nearest.get(0)
    if black is None:
        black = _nearest(0, entries)
    indices[len(tiles):] = bytes([black]) * (TILES_PER_SCREEN - len(tiles))
    return bytes(indices)