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
    DLX3 only: the v7 LITERAL SPAN section (tools/encoder/spans.py) --
      u16 nspans, then per span { u32 GVRAM address, u16 coarse displacement,
      c*48 B pixels, u16 fine displacement, f*4 B pixels }
    then payloads in block order: V1 -> 1 byte, V4 -> 4 bytes, RAW -> 16 bytes

  The span section is between the header and the block payload, not after it,
  because the 68000 has to reach it without first parsing something of variable
  length: the mode header is a fixed 768 bytes, so the section starts at a known
  offset and the block payload starts wherever the span walk finishes.  Every
  span record is a multiple of 4 bytes long, so nothing inside needs padding.

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, spans as SP

# 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 build_records(m, enc, span_mode):
    """The per-frame records of the container, in order.

    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.

    Factored out of main() so tools/analysis/16_span_roundtrip.py can build the
    same bytes the shipping encoder does -- a round-trip gate that rebuilt the
    records itself would be testing its own copy of the format.
    """
    nbx = m["W"] // 4
    out = []
    for f, im in enumerate(m["idx"]):
        mode = enc["modes"][f]
        sp = enc.get("spans", [[]] * len(m["idx"]))[f]
        rec = (pack_modes(mode)
               + (SP.serialise(sp) if span_mode else b"")
               + 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 reported is fiction
        assert len(rec) == enc["sizes"][f], (f, len(rec), enc["sizes"][f])
        out.append(rec)
    return out


def write_container(path, m, frames, fps, k1, k4, span_mode):
    """Write the whole container.  Returns (total bytes, video bytes, pad)."""
    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"DLX3" if span_mode else b"DLX2")
           + struct.pack(">HHHHHH", m["W"], m["H"], fps, len(frames), k1, k4)
           + struct.pack(">IIII", off_pal, off_cb1, off_cb4, off_frm))
    assert len(hdr) == 32, len(hdr)

    frm_pad = 0
    with open(path, "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, rec in enumerate(frames):
            fh.write(struct.pack(">I", len(rec))); fh.write(rec)
            if i + 1 < len(frames):          # nothing follows the last record
                n = -(4 + len(rec)) % 4
                fh.write(b"\0" * n); frm_pad += n
    total = os.path.getsize(path)
    return total, sum(len(r) + 4 for r in frames) + frm_pad, frm_pad


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("--kbps", type=float, default=None,
                    help="override the profile's bitrate CEILING. The profile "
                         "is a rate point on a delivery medium; this is for "
                         "asking what the codec does at another one -- e.g. "
                         "the 488 KB/s bus figure the span analyses of "
                         "FINDINGS 30/40 are scored against.")
    ap.add_argument("--span-kbps", type=float, default=None,
                    help="byte ceiling the SPAN pass may draw on, if it differs "
                         "from the profile's. The profile is a quality rate "
                         "point; the pipe is hardware. Bytes between the two "
                         "buy a better picture if spent on lam and the 68000's "
                         "deadline if spent on spans -- and nothing at all if "
                         "left unspent (FINDINGS 41.2). 488 is the bus figure "
                         "tools/analysis/14_dmac_chain.py scores against.")
    ap.add_argument("--spans", choices=("off", "need", "all"), default="need",
                    help="v7 literal spans (FINDINGS 40). `need` (default) "
                         "spends container bytes on spans only where a frame "
                         "misses the 68000's decode deadline; `all` spends "
                         "every profitable byte, which is the model "
                         "14_dmac_chain.py scores; `off` emits DLX2.")
    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 = dict(RC.PROFILES[a.profile])
    if a.kbps is not None:
        prof["kbps"] = a.kbps
        prof["desc"] = f"{prof['desc']} -- bitrate overridden to {a.kbps:g} KB/s"
    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
    span_mode = None if (a.spans == "off" or not rc) else a.spans

    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,
                                        span_mode=span_mode,
                                        span_kbps=a.span_kbps)
    else:
        # Spans are a rate-control-era mode: `need` has no meaning without a
        # per-frame byte allowance to spend, so --fixed-lam emits DLX2.
        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"]

    frames = build_records(m, enc, span_mode)
    nspans = sum(len(x) for x in enc.get("spans", []))

    total, vid, frm_pad = write_container(a.out, m, frames, a.fps, k1, k4,
                                          span_mode)
    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 span_mode:
        spf = np.array([len(x) for x in enc["spans"]])
        spb = np.array([SP.section_bytes(x) for x in enc["spans"]])
        # a run of L blocks is four spans of 4L pixels, so a block is 16 span
        # pixels -- not 4, which would count each block four times over
        blk = np.array([sum(len(p) for _, _, p in x) // 16 for x in enc["spans"]])
        print(f"  v7 spans ({span_mode}): {nspans:,} over {len(idx)} frames, "
              f"median {np.median(spf):.0f}/frame, max {spf.max()}/frame; "
              f"{100*np.mean(spb)/np.mean([len(p) for p in frames]):.1f}% of the "
              f"container")
        print(f"    frames with any span: {int((spf>0).sum())}/{len(idx)}; "
              f"blocks painted by one: median {np.median(blk):.0f}, "
              f"max {blk.max()} of {m['nb']} "
              f"({100*blk.max()/m['nb']:.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"]])
    # enc["cycles"] already carries the span PAINTING clocks; H.cycles() sees
    # only the mode map, in which a spanned block reads SKIP, so re-deriving
    # here would report a frame as fitting on the strength of work the encoder
    # moved into the span section rather than removed.
    cyc = (np.asarray(enc["cycles"]) if "cycles" in enc
           else 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()