Dragon-s-Lair-X68k/tools/encoder/vq_hybrid.py

#!/usr/bin/env python3
"""Cinepak-style hybrid VQ with a RAW escape: per-4x4-block choice of
SKIP / V1 (one 4x4 codeword) / V4 (four 2x2 codewords) / RAW (16 literal indices).

Flat 4x4 VQ at k=256 visibly destroys Bluth's ink linework (see docs/FINDINGS.md).
The standard fix is to let detailed blocks spend 4x the bits. Rate control picks
the split per block by rate-distortion, so the bitrate ceiling stays deterministic
-- which is the whole reason we chose VQ over a lossless delta.

The RAW mode is what makes ONE codec serve both shipping targets (session 2
user decision: SASI and SCSI quality modes). As lam -> 0 the encoder buys RAW
blocks until the frame is pixel-exact against the palettised source, so the
SCSI profile is not a second codec -- it is the same bitstream with the rate
knob opened up. The 68000 decoder needs no extra path: RAW is a straight copy,
which is cheaper than V4.

Bitstream per frame (what the 68000 actually parses):
  2 bits/block header, packed:  00=SKIP  01=V1  10=V4  11=RAW
  then the payload in block order: V1 -> 1 index, V4 -> 4, RAW -> 16

STRUCTURE (session 6). The encoder is FRAME-DRIVABLE: `frame_ctx` / `decide` /
`paint` expose one frame at a time so a caller can choose `lam` per frame and
feed back the frame it actually emitted. That is not a convenience -- it is the
fix for FINDINGS 26.1. This codec is temporally recursive (SKIP copies the
previous RECONSTRUCTION), so any rate control that picks frames out of
independently-encoded whole-sequence runs desynchronises the encoder from the
decoder. `encode()` is now a thin loop over the per-frame API and stays the
fixed-lam path.

The split is also what makes rate control affordable: `l1`/`l4g` and their
errors depend on neither `lam` nor `prev`, so they are computed once per frame
and a lam search only re-runs the argmin.
"""
import numpy as np, sys
from PIL import Image
import vq as VQ

LUMA = VQ.LUMA

# byte cost of each mode's payload, per block. The 2-bit header is paid by
# every block regardless, so it drops out of the mode comparison.
_HDR_BYTES_PER_BLOCK = 2 / 8.0
RAW_BYTES = 16.0                  # literal palette bytes, never indices

# ---------------------------------------------------------------------------
# CYCLE cost of each mode, per block, MEASURED on the 68000 (FINDINGS 28.2,
# tools/bench/decode.lua). This is the other axis: `lam` prices bytes, `mu`
# prices cycles, and the two are not proportional -- V4 is 4x a V1 block in
# bytes and 1.49x in cycles.
#
# SKIP IS NOT A CONSTANT, and it is the one trap in here. A SKIP block costs
# 13.25 cycles when all four blocks sharing its header byte are SKIP (one
# `tst.b` clears the group) and ~45 when it sits in a mixed byte -- so its
# price depends on its NEIGHBOURS, which a per-block lagrangian cannot see.
# The way out is that the two uses do not need the same number:
#   * `decide` uses C_SKIP_RANK purely to RANK modes within a block. SKIP is
#     the cheapest mode either way, so the choice only scales the incentive:
#     the V1-SKIP gap moves 12% between the two candidates.
#   * `cycles()` scores a WHOLE frame with the exact clustered rule, and that
#     is what the rate controller bisects against. Nothing downstream of the
#     mode decision uses the ranking constant.
C_V1, C_V4, C_RAW = 299.9, 448.2, 400.4
C_SKIP_CLUSTERED = 53.0 / 4       # all-SKIP header byte: one tst.b for four
C_SKIP_MIXED = 45.0               # a SKIP block inside a mixed byte
C_SKIP_RANK = C_SKIP_CLUSTERED    # ranking only -- see above
MODE_CYCLES = np.array([C_SKIP_RANK, C_V1, C_V4, C_RAW], dtype=np.float64)

FRAME_CYCLES_12FPS = 10_000_000 / 12.0     # 833,333, x68k.cpp:1133


def cycles(mode):
    """Exact decode cost of one frame's mode map, in 68000 cycles.

    Single source of truth: tools/analysis/11_cpu_budget.py imports this, and
    it reproduces the four frames timed on the 68000 to within 1 point
    (FINDINGS 28.2). Instruction cycles against zero-wait-state memory, so a
    LOWER BOUND like every 68000 figure since FINDINGS 24."""
    g = np.asarray(mode).reshape(-1, 4)        # one header byte = four blocks
    allskip = (g == 0).all(1)
    c = allskip.sum() * 4 * C_SKIP_CLUSTERED
    mm = g[~allskip]
    c += (mm == 0).sum() * C_SKIP_MIXED
    c += (mm == 1).sum() * C_V1
    c += (mm == 2).sum() * C_V4
    c += (mm == 3).sum() * C_RAW
    return float(c)


def blocks_of(idx, pal, bw, bh):
    return VQ.blockify(idx, pal, bw, bh)


def build(frames_dir, k1=256, k4=256, iters=16, lam=0.0):
    rgb = VQ.load_frames(frames_dir)
    H, W = rgb[0].shape[:2]
    ref, pal = VQ.scene_palette(rgb)
    idx = VQ.palettise(rgb, ref)

    # --- two codebooks, trained on the whole scene ---
    X1 = np.concatenate([blocks_of(i, pal, 4, 4) for i in idx])
    C1, _ = VQ.kmeans(X1, k1, iters)
    cb1 = VQ.snap_codebook(C1, pal, 4, 4)                  # (k1,16) palette idx
    C1s = (pal[cb1].astype(np.float32) * LUMA).reshape(k1, -1)

    X4 = np.concatenate([blocks_of(i, pal, 2, 2) for i in idx])
    C4, _ = VQ.kmeans(X4, k4, iters)
    cb4 = VQ.snap_codebook(C4, pal, 2, 2)                  # (k4,4) palette idx
    C4s = (pal[cb4].astype(np.float32) * LUMA).reshape(k4, -1)
    return dict(rgb=rgb, pal=pal, idx=idx, H=H, W=W,
                cb1=cb1, C1s=C1s, cb4=cb4, C4s=C4s, k1=k1, k4=k4,
                nbx=W // 4, nby=H // 4, nb=(W // 4) * (H // 4),
                palw=pal.astype(np.float32) * LUMA)


def _v1_recon(lab1, cb1, H, W):
    return VQ.unblockify(lab1, cb1, H, W, 4, 4)


# --- block <-> image reshapes (no copy where numpy can avoid one) -------------

def to_blocks(a, nbx, nby):
    """(H,W) -> (nb,4,4) in block raster order."""
    return a.reshape(nby, 4, nbx, 4).transpose(0, 2, 1, 3).reshape(-1, 4, 4)


def from_blocks(b, nbx, nby):
    """(nb,4,4) -> (H,W)."""
    return b.reshape(nby, nbx, 4, 4).transpose(0, 2, 1, 3).reshape(nby * 4, nbx * 4)


def default_idx_bytes(m):
    """Size of ONE codebook index in the bitstream. k>256 needs 2 bytes, which
    doubles what V1 and V4 actually cost -- an RD model that ignores that
    systematically over-picks V4 and under-reports the bitrate (FINDINGS 14)."""
    return 1 if max(m["k1"], m["k4"]) <= 256 else 2


# --- per-frame API -----------------------------------------------------------

def frame_symbols(m, f):
    """lam- and prev-INDEPENDENT part of a frame: codeword assignments and
    their errors.

    Cached, because a lam search re-uses them unchanged and `VQ.assign` is the
    expensive call in the encoder -- 22.8 of 24.6 ms per frame, measured. That
    cache is what makes per-frame rate control affordable: a 12-step search
    over 120 frames costs 0.3 s, against 49 s for the equivalent done by
    re-running whole-sequence encodes.

    The cache holds ONE frame. Every caller works a frame at a time, and at
    ~133 KB of intermediates per frame a whole-sequence cache would cost
    900 MB on a 9.4-minute stream for no benefit."""
    cache = m.get("_sym")
    if cache is not None and cache[0] == f:
        return cache[1]
    pal, W, nb = m["pal"], m["W"], m["nb"]
    im = m["idx"][f]

    B1 = blocks_of(im, pal, 4, 4)                       # (nb,48)
    l1 = VQ.assign(B1, m["C1s"])
    e1 = ((B1 - m["C1s"][l1]) ** 2).sum(1)

    B4 = blocks_of(im, pal, 2, 2)                       # (nb*4,12) in 2x2 raster
    l4 = VQ.assign(B4, m["C4s"])
    e4raw = ((B4 - m["C4s"][l4]) ** 2).sum(1)
    # regroup 2x2 blocks into their parent 4x4 block
    q = _group_2x2_into_4x4(np.arange(nb * 4), W)
    e4 = e4raw[q].reshape(nb, 4).sum(1)
    l4g = l4[q].reshape(nb, 4)

    s = dict(l1=l1, e1=e1, l4g=l4g, e4=e4,
             src_blocks=to_blocks(im, m["nbx"], m["nby"]))
    m["_sym"] = (f, s)
    return s


def frame_ctx(m, f, prev, idx_bytes=None):
    """Everything needed to decide one frame at any lam, given the frame that
    will actually precede it in the emitted stream."""
    s = frame_symbols(m, f)
    nb = m["nb"]
    if prev is None:
        eS = np.full(nb, np.inf)
        prev_blocks = None
    else:
        # SKIP distortion = this frame against the previous RECONSTRUCTION,
        # in the same luma-weighted space the codebooks were trained in.
        d = ((m["palw"][m["idx"][f]] - m["palw"][prev]) ** 2).sum(2)
        eS = d.reshape(m["nby"], 4, m["nbx"], 4).sum((1, 3)).ravel()
        prev_blocks = to_blocks(prev, m["nbx"], m["nby"])
    return dict(f=f, sym=s, eS=eS, prev_blocks=prev_blocks, nb=nb,
                idx_bytes=default_idx_bytes(m) if idx_bytes is None else idx_bytes)


def decide(ctx, lam, mu=0.0):
    """Lagrangian mode decision at one lam and one mu. Returns (mode, bytes).

    Minimises `distortion + lam*bytes + mu*cycles` per block. `mu=0` is the
    byte-only decision every session before 8 made; the machine's binding
    budget is cycles, and bytes and cycles do not rank the modes the same way
    (V4 is 4x V1 in bytes, 1.49x in cycles; RAW is dearer than V4 in bytes and
    CHEAPER in cycles, so mu inverts that preference -- FINDINGS 28.8).

    Cheap by design: no painting, no image-sized work. A search calls this a
    dozen times per lam step and paints once."""
    ib = ctx["idx_bytes"]
    s = ctx["sym"]
    mc = mu * MODE_CYCLES
    cost = np.stack([ctx["eS"] + mc[0],
                     s["e1"] + lam * (1.0 * ib) + mc[1],
                     s["e4"] + lam * (4.0 * ib) + mc[2],
                     np.full(ctx["nb"], lam * RAW_BYTES + mc[3])])
    mode = np.argmin(cost, axis=0).astype(np.uint8)
    return mode, frame_bytes(mode, ctx["nb"], ib)


def frame_bytes(mode, nb, idx_bytes):
    nV1 = int((mode == 1).sum()); nV4 = int((mode == 2).sum())
    nR = int((mode == 3).sum())
    return (nb * _HDR_BYTES_PER_BLOCK
            + (nV1 + nV4 * 4) * idx_bytes + nR * RAW_BYTES)


def paint(m, ctx, mode):
    """Reconstruct the frame the decoder will produce for this mode map."""
    nbx, nby = m["nbx"], m["nby"]
    s = ctx["sym"]
    ob = np.empty((ctx["nb"], 4, 4), dtype=np.uint8)
    sel = mode == 0
    if sel.any():
        ob[sel] = ctx["prev_blocks"][sel]
    sel = mode == 1
    if sel.any():
        ob[sel] = m["cb1"][s["l1"][sel]].reshape(-1, 4, 4)
    sel = mode == 2
    if sel.any():
        # (n, sub_y, sub_x, py, px) -> (n, sub_y, py, sub_x, px) -> (n,4,4)
        c = m["cb4"][s["l4g"][sel]].reshape(-1, 2, 2, 2, 2)
        ob[sel] = c.transpose(0, 1, 3, 2, 4).reshape(-1, 4, 4)
    sel = mode == 3
    if sel.any():
        ob[sel] = s["src_blocks"][sel]
    return from_blocks(ob, nbx, nby)


def encode_frame(m, f, prev, lam, idx_bytes=None, mu=0.0):
    """One frame at one lam against one previous reconstruction."""
    ctx = frame_ctx(m, f, prev, idx_bytes)
    mode, sz = decide(ctx, lam, mu)
    return dict(recon=paint(m, ctx, mode), mode=mode, size=sz,
                l1=ctx["sym"]["l1"], l4g=ctx["sym"]["l4g"], ctx=ctx)


def encode(m, lam=0.02, skip_thresh=0.0, idx_bytes=None):
    """Fixed-lam whole-sequence encode: a loop over the per-frame API.

    lam = lagrangian rate weight (bytes -> squared-error units).
    Higher lam => more V1/SKIP => smaller & softer.

    For a rate-controlled encode use ratectl.encode_rate_controlled(), which
    drives the same per-frame API and varies lam. Do NOT reassemble a sequence
    out of several fixed-lam runs of this function -- FINDINGS 26.1."""
    if idx_bytes is None:
        idx_bytes = default_idx_bytes(m)
    recon, modes, sizes, l1s, l4gs = [], [], [], [], []
    prev = None
    for f in range(len(m["idx"])):
        r = encode_frame(m, f, prev, lam, idx_bytes)
        recon.append(r["recon"]); modes.append(r["mode"]); sizes.append(r["size"])
        l1s.append(r["l1"]); l4gs.append(r["l4g"])
        prev = r["recon"]
    return dict(recon=recon, modes=modes, sizes=np.array(sizes),
                l1=l1s, l4g=l4gs, nb=m["nb"])


def _group_2x2_into_4x4(a, W):
    """map 2x2-block raster order -> (nb4, 4) grouping by parent 4x4 block"""
    n2x = W // 2
    n2y = len(a) // n2x
    g = a.reshape(n2y, n2x)
    g = g.reshape(n2y // 2, 2, n2x // 2, 2).transpose(0, 2, 1, 3)
    return g.reshape(-1)


def evaluate(m, enc, fps=12):
    pal = m["pal"]
    rec = [pal[i] for i in enc["recon"]]
    src = [pal[i] for i in m["idx"]]
    p_vq = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], rec)])
    p_pal = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], src)])
    mo = np.concatenate(enc["modes"])
    sz = enc["sizes"].mean()
    return dict(psnr=p_vq, pal=p_pal, loss=p_pal - p_vq, bytes=sz,
                kbps=sz * fps / 1024,
                skip=100 * (mo == 0).mean(), v1=100 * (mo == 1).mean(),
                v4=100 * (mo == 2).mean(), raw=100 * (mo == 3).mean())