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

#!/usr/bin/env python3
"""Encode one scene to the DLX bitstream, at a chosen quality profile.

    python3 tools/encoder/encode.py <frames_dir> <out.dlx> [--profile scsi]
                                    [--lam N] [--fps 12] [--preview out.png]
                                    [--fixed-lam] [--rc-floor profile|open]

Rate control is ON by default: lam is bisected per frame under a leaky bucket
so the profile's bitrate is a ceiling rather than an average hope. `--fixed-lam`
restores session 5's behaviour, which overshoots by 18-34% on sustained action
(FINDINGS 25.3). `--rc-floor` picks the quality floor: `profile` (default) never
spends more than the fixed-lam profile would, so it can only ever help; `open`
lets quiet frames spend the whole allowance and lands the mean ON target.

Container (little-endian is WRONG here -- the 68000 is big-endian, so every
multi-byte field is big-endian and the decoder can read it with a plain move.w):

  header, 32 bytes
    0   'DLX2'                 magic ('DLX1' = the same, unaligned; still read)
    4   u16 width, u16 height
    8   u16 fps,   u16 nframes
    12  u16 k1,    u16 k4      codebook sizes
    16  u32 palette offset     (256 * 3 bytes, RGB888 -- the player converts
                                to the X68000's GRB555 at load time)
    20  u32 cb1 offset         (k1 * 16 bytes of palette indices)
    24  u32 cb4 offset         (k4 *  4 bytes)
    28  u32 frames offset
  then, per frame, each record starting on a 4-BYTE BOUNDARY (0-3 zero pad
  bytes before it; a 68000 takes an address error, not a slow read, on an odd
  `move.l` -- FINDINGS 28.3):
    u32 payload length, then
    ceil(nblocks*2/8) bytes of 2-bit mode headers, MSB-first, block raster order
    then payloads in block order: V1 -> 1 byte, V4 -> 4 bytes, RAW -> 16 bytes

Codebooks are emitted as palette INDICES, not pixels. The player expands them
once at load time into word-per-pixel form so the blitter can movem them
straight into GVRAM -- k1=1024 costs 1024*16*2 = 32 KB of the 2 MB.
"""
import argparse, struct, sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
import vq as VQ, vq_hybrid as H, ratectl as RC

# Measured on the emulated 68000, FINDINGS 24. Instruction cycles against
# zero-wait-state memory, so these are floors, not hardware predictions.
BLIT_PCT      = 53.6      # V1: compose in RAM, then a row-linear movem.l blit
DIRECT_PCT    = 76.6      # V4: write every block straight into GVRAM
CROSSOVER_PCT = 100 * BLIT_PCT / DIRECT_PCT


def pack_modes(mode):
    """2 bits per block, MSB-first -- cheap for the 68000 to shift out."""
    n = len(mode)
    out = bytearray((n * 2 + 7) // 8)
    for i, m in enumerate(mode):
        out[i // 4] |= (int(m) & 3) << (6 - 2 * (i % 4))
    return bytes(out)


def frame_payload(mode, l1, l4g, src_idx, nbx):
    body = bytearray()
    for b, mo in enumerate(mode):
        if mo == 1:
            body += _idx(l1[b])
        elif mo == 2:
            for j in range(4):
                body += _idx(l4g[b][j])
        elif mo == 3:
            by, bx = divmod(b, nbx)
            body += src_idx[by*4:by*4+4, bx*4:bx*4+4].tobytes()
    return bytes(body)


def _idx(v):
    """codebook index: 1 byte if it fits, else big-endian u16.
    k>256 means 2-byte indices -- decided once by the header, not per block."""
    v = int(v)
    return bytes([v]) if _IDX_BYTES == 1 else struct.pack(">H", v)


_IDX_BYTES = 1


def main():
    global _IDX_BYTES
    ap = argparse.ArgumentParser()
    ap.add_argument("frames_dir"); ap.add_argument("out")
    ap.add_argument("--profile", choices=list(RC.PROFILES), default="scsi")
    ap.add_argument("--lam", type=float, default=None)
    ap.add_argument("--fps", type=int, default=12)
    ap.add_argument("--iters", type=int, default=16)
    ap.add_argument("--fixed-lam", action="store_true",
                    help="disable rate control (session 5 behaviour)")
    ap.add_argument("--rc-floor", choices=("profile", "open"), default="profile",
                    help="quality floor for rate control")
    ap.add_argument("--bucket-frames", type=int, default=8,
                    help="leaky-bucket depth, in frame budgets")
    ap.add_argument("--no-cpu-fit", action="store_true",
                    help="drop the per-frame 68000 decode ceiling (session 7 "
                         "behaviour: 31%% of frames on hard content do not fit)")
    ap.add_argument("--prefill", type=float, default=0.0,
                    help="how full the player's buffer is assumed to be at "
                         "scene start, as a fraction of the bucket (0 = cold "
                         "buffer after a seek, the conservative assumption)")
    ap.add_argument("--preview")
    a = ap.parse_args()

    prof = RC.PROFILES[a.profile]
    lam = a.lam if a.lam is not None else prof["lam"]
    k1, k4 = prof["k1"], prof["k4"]
    _IDX_BYTES = 1 if max(k1, k4) <= 256 else 2

    # An explicit --lam is a request for that lam, so it implies --fixed-lam.
    rc = not (a.fixed_lam or a.lam is not None)
    lam_lo = lam if a.rc_floor == "profile" else 1.0
    # The CPU ceiling is hardware, not taste: without it 31%% of frames on the
    # worst sustained window do not decode in time on a stock 68000, and with
    # it that is one frame -- the intra frame -- for 0.26 dB. FINDINGS 31.
    cyc_budget = None if a.no_cpu_fit else RC.FRAME_CYCLES

    print(f"profile {a.profile}: {prof['desc']}")
    if rc:
        print(f"  target {prof['kbps']} KB/s CEILING, rate-controlled: "
              f"lam bisected per frame in [{lam_lo:g}, {RC.LAM_CLIFF:g}], "
              f"{a.bucket_frames}-frame bucket")
        print(f"  CPU ceiling: " + (f"mu bisected per frame against "
              f"{RC.FRAME_CYCLES:,.0f} cycles (12fps, stock 68000)"
              if cyc_budget else "OFF (--no-cpu-fit)"))
    else:
        print(f"  target {prof['kbps']} KB/s, FIXED lam={lam} (no rate control)")
    print(f"  k1={k1} k4={k4}, {_IDX_BYTES}-byte indices")

    m = H.build(a.frames_dir, k1=k1, k4=k4, iters=a.iters)
    if rc:
        enc = RC.encode_rate_controlled(m, prof["kbps"], fps=a.fps,
                                        bucket_frames=a.bucket_frames,
                                        lam_lo=lam_lo, prefill=a.prefill,
                                        cycle_budget=cyc_budget)
    else:
        enc = H.encode(m, lam=lam)
    r = H.evaluate(m, enc, fps=a.fps)

    H_, W_ = m["H"], m["W"]; nbx = W_ // 4
    pal, idx = m["pal"], m["idx"]

    # The encoder hands back the symbols it actually chose. Re-deriving them
    # here (as session 5 did) is a second chance to disagree with the encoder,
    # and with per-frame rate control the mode map is no longer reproducible
    # from a single lam anyway.
    frames = []
    for f, im in enumerate(idx):
        mode = enc["modes"][f]
        frames.append(pack_modes(mode)
                      + frame_payload(mode, enc["l1"][f], enc["l4g"][f], im, nbx))
        # the rate controller budgets exactly these bytes -- if that ever drifts
        # from the container, every bitrate figure below is fiction
        assert len(frames[-1]) == enc["sizes"][f], (f, len(frames[-1]), enc["sizes"][f])

    palette = m["pal"][:256]
    if len(palette) < 256:
        palette = np.vstack([palette, np.zeros((256 - len(palette), 3), np.uint8)])
    pal_b = palette.astype(np.uint8).tobytes()
    cb1_b = m["cb1"].astype(np.uint8).tobytes()
    cb4_b = m["cb4"].astype(np.uint8).tobytes()

    off_pal = 32
    off_cb1 = off_pal + len(pal_b)
    off_cb4 = off_cb1 + len(cb1_b)
    off_frm = off_cb4 + len(cb4_b)
    # DLX2: every frame record starts on a 4-byte boundary, including the
    # first. Payload lengths are arbitrary, so end-to-end records land on odd
    # addresses -- and `move.l (a0)+` at an odd address is an ADDRESS ERROR on
    # a 68000, not a slow read. It vectors into the IPL and looks exactly like
    # an infinite loop (FINDINGS 28.3). tools/bench/prep_dlx.py has been
    # realigning at load time; the container now carries it.
    tbl_pad = -off_frm % 4
    off_frm += tbl_pad
    hdr = (b"DLX2" + struct.pack(">HHHHHH", W_, H_, a.fps, len(idx), k1, k4)
           + struct.pack(">IIII", off_pal, off_cb1, off_cb4, off_frm))
    assert len(hdr) == 32, len(hdr)

    frm_pad = 0
    with open(a.out, "wb") as fh:
        fh.write(hdr); fh.write(pal_b); fh.write(cb1_b); fh.write(cb4_b)
        fh.write(b"\0" * tbl_pad)
        for i, p in enumerate(frames):
            fh.write(struct.pack(">I", len(p))); fh.write(p)
            if i + 1 < len(frames):          # nothing follows the last record
                n = -(4 + len(p)) % 4
                fh.write(b"\0" * n); frm_pad += n

    total = os.path.getsize(a.out)
    vid = sum(len(p) + 4 for p in frames) + frm_pad
    print(f"  wrote {a.out}: {total} B  "
          f"(header+tables {total-vid} B, video {vid} B)")
    print(f"  DLX2 4-byte record alignment: {frm_pad} B over {len(frames)} frames "
          f"({frm_pad/len(frames):.2f} B/frame = {frm_pad/len(frames)*a.fps:.0f} B/s)")
    print(f"  {vid/len(idx):.0f} B/frame -> {vid/len(idx)*a.fps/1024:.1f} KB/s video"
          f" + {RC.AUDIO_KBPS} KB/s audio = {vid/len(idx)*a.fps/1024+RC.AUDIO_KBPS:.1f} KB/s")
    print(f"  PSNR {r['psnr']:.2f} dB   palette ceiling {r['pal']:.2f} dB   "
          f"loss {r['loss']:.2f} dB")
    print(f"  modes: SKIP {r['skip']:.1f}%  V1 {r['v1']:.1f}%  "
          f"V4 {r['v4']:.1f}%  RAW {r['raw']:.1f}%")
    if rc:
        rr = RC.summarise(m, enc, prof["kbps"], fps=a.fps)
        lm = enc["lam"]
        print(f"  rate control: per-frame budget {enc['budget']:.0f} B, "
              f"bucket {enc['cap']:.0f} B ({a.bucket_frames} frames), "
              f"prefill {100*a.prefill:.0f}%")
        print(f"    lam: min {lm.min():.1f}  median {rr['lam_med']:.1f}  "
              f"p90 {rr['lam_p90']:.1f}  max {rr['lam_max']:.1f}")
        print(f"    frames over the per-frame budget (banked by the bucket): "
              f"{rr['over']:.0f}%")
        print(f"    frames that could not fit even at the lam={RC.LAM_CLIFF:g} "
              f"cliff: {rr['overrun']}/{len(lm)}")

    # PER-FRAME DECODE COST, from the measured per-mode block costs
    # (FINDINGS 28.2, vq_hybrid.cycles). The mean cannot answer this: a scene
    # cut is ~100% non-SKIP and a held frame near 0%, so their mean describes
    # no real frame. What matters is how many frames MISS, and by how much.
    #
    # This replaces the per-frame blit-vs-direct path choice that used to be
    # printed here. That plan is withdrawn -- mixing the two paths displays
    # stale pixels on 70 of 120 frames, and there was never a crossover to
    # begin with, because the compose path pays the blit ON TOP of decoding.
    # FINDINGS 28.1/28.4. The player has one path and no reference frame.
    ns = np.array([100 * (mm != 0).mean() for mm in enc["modes"]])
    cyc = np.array([H.cycles(mm) for mm in enc["modes"]])
    pct = 100 * cyc / RC.FRAME_CYCLES
    miss = int((pct > 100).sum())
    print(f"  non-SKIP blocks/frame: median {np.median(ns):.1f}%  "
          f"p90 {np.percentile(ns, 90):.1f}%  max {ns.max():.1f}%")
    print(f"  decode cost: median {np.median(pct):.1f}%  "
          f"p90 {np.percentile(pct, 90):.1f}%  max {pct.max():.1f}% "
          f"of a {a.fps}fps frame")
    print(f"  frames that do NOT decode in time: {miss}/{len(pct)} "
          f"({100*miss/len(pct):.0f}%)"
          + (f"  -- worst {pct.max():.1f}%" if miss else ""))
    if rc and cyc_budget:
        rr2 = RC.summarise(m, enc, prof["kbps"], fps=a.fps)
        print(f"    mu: median {rr2['mu_med']:.4f}  max {rr2['mu_max']:.3f}   "
              f"frames needing any mu at all: {int((enc['mu'] > 0).sum())}/{len(pct)}")
        print(f"    frames that cannot fit even at mu={RC.MU_CLIFF:g} "
              f"(emitted late on purpose): {rr2['late']}")

    if a.preview:
        from PIL import Image
        f = len(idx) // 2
        gap = np.full((H_ * 3, 4, 3), 40, np.uint8)
        st = np.concatenate([VQ.zoom(m["rgb"][f], 3), gap,
                             VQ.zoom(pal[idx[f]], 3), gap,
                             VQ.zoom(pal[enc["recon"][f]], 3)], axis=1)
        Image.fromarray(st).save(a.preview)
        print(f"  preview -> {a.preview}  (source | palette ceiling | decoded)")


if __name__ == "__main__":
    main()