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

#!/usr/bin/env python3
"""Rate control: hit a target bitrate exactly, so one encoder serves both targets.

USER DECISION (session 2): ship TWO quality modes, SASI and SCSI. The codec's
bitrate ceiling is a build parameter; the encoder is otherwise identical.

Mechanism: the hybrid encoder's lagrangian `lam` trades distortion for bytes
monotonically, so per frame we binary-search lam to land inside a byte budget.
A leaky bucket lets a quiet frame bank bytes that an action frame can spend --
without that, quiet frames waste budget and action frames stay ugly.

The ceiling is HARD: the 68000 streams at a fixed rate off the disk, and a frame
that overruns is a dropped frame, not a slow frame.

STATUS, session 6: WIRED IN and sound. `encode.py` rate-controls by default
for a profile; `--fixed-lam` restores the old behaviour. The lam-ladder of
session 5 was replaced by a per-frame bisection that drives the encoder one
frame at a time and feeds back the frame it actually emitted -- see
encode_rate_controlled(), and FINDINGS 26 for why the ladder could not be
fixed by tuning. Regression test: tools/analysis/09_ratectl_drift.py.
"""
import numpy as np
import vq_hybrid as H

# Profiles. Bandwidths are the sustained-read figures the player can rely on;
# see docs/FINDINGS.md 5 -- these are FOLKLORE-grade until the disk benchmark
# is unblocked, so they are deliberately conservative fractions of the quoted
# ceiling (audio, seeks and container overhead come out of the same pipe).
# Calibrated against the corrected rate-distortion curve (FINDINGS 14-15) AND
# against the bus and CPU budgets (FINDINGS 17). Session 2 initially set these
# far too low: 45 / 75 KB/s is 12% of the pessimistic SASI folklore figure and
# 7% of the SCSI one. Nothing justified that -- the numbers were read off the
# knee of the RD curve, not off the hardware.
#
# What actually bounds the high end:
#   - Bus: unmeasured. ~300-500 KB/s SASI / ~1 MB/s SCSI, both FOLKLORE.
#     This is the binding unknown and the reason the disk benchmark matters.
#   - CPU: a full-frame blit is **53.6%** of the 12fps budget -- MEASURED on the
#     emulated 68000, session 5, FINDINGS 24. This line previously said 38%,
#     which was an estimate and was wrong by 41%. And 53.6% is a floor: MAME
#     models no GVRAM wait states, so real hardware is worse.
#     "Raising the bitrate is nearly free on CPU" survives but is now much
#     tighter. It rests on RAW being the cheapest mode to blit, which is still
#     true, but the display path alone now eats over half the frame before any
#     decoding happens.  The per-frame path choice of FINDINGS 25.6 (blit vs
#     direct-to-GVRAM, whichever is cheaper for that frame) brings the median
#     back to ~37% and caps the worst case at 53.6%.
#   - Entropy coding is NOT the way to buy headroom here: deflate decode is
#     ~216% of the frame budget on a 68000 and even LZ4 is ~54%. See FINDINGS 17.
#     The rates below are therefore RAW payload, no entropy coding.
#
# The two profiles are the SAME codec, decoder and bitstream -- only `lam` differs.
# `lam` here is a FLOOR, not a setting: encode.py rate-controls by default and
# bisects lam per frame in [lam, LAM_CLIFF] to keep under `kbps`. The floor is
# what a quiet frame is allowed to spend, so rate control can only ever spend
# less than session 5's fixed-lam encoder did. FINDINGS 27.
PROFILES = {
    "sasi": dict(kbps=110, lam=60.0, k1=256, k4=256,
                 desc="stock 10MHz ACE/EXPERT, SASI",
                 quality="36.9 dB on 00020 / 29.6 dB on 00146 / 27.2 dB on the "
                         "Singe window at 109.5 KB/s (session 5's fixed lam "
                         "gave 27.8 dB there, but at 137.4 KB/s)",
                 util="~105 KB/s = 35% of the pessimistic 300 KB/s SASI figure"),
    "scsi": dict(kbps=280, lam=10.0, k1=256, k4=256,
                 desc="Super/XVI, or CZ-6BS1 board in a 10MHz machine",
                 quality="39.4 dB on 00020 / 32.3 dB on 00146 / 29.9 dB on the "
                         "Singe window at 280.0 KB/s (session 5's fixed lam "
                         "gave 30.8 dB there, but at 381.6 KB/s)",
                 util="~275 KB/s = 28% of the 1 MB/s SCSI folklore figure"),
}
# lam=0 is PIXEL-EXACT against the palettised frame (0.00 dB loss) at ~450 KB/s
# of raw payload, and costs 53.6% of the CPU budget (not the 38% written here
# before session 5 -- FINDINGS 24). If the blocked disk
# benchmark confirms SCSI sustains >=800 KB/s, the `scsi` profile should become
# lam=0 and the port ships transparent video. That decision is waiting on a
# measurement, not on a design choice.

# Hard ceiling on the rate-control search. FINDINGS 15 puts the quality cliff
# between lam=800 and lam=2000. Above it a frame has not been rate-controlled,
# it has been destroyed, so the search stops here and lets the frame overrun
# instead (FINDINGS 26.2). The old ladder ran to lam=2e5, 250x past shippable.
LAM_CLIFF = 800.0

# Ceiling on the CYCLE search. mu prices a cycle in the same units lam prices a
# byte, so the scale that matters is set by their ratio: at the `sasi` floor of
# lam=60, mu=0.2 makes a V1 block's 300 cycles cost what its 1 payload byte
# costs. MU_CLIFF=100 is three decades past that: a V1 block priced at 30,000
# distortion units.
#
# It does NOT freeze the picture, and that is the point. At MU_CLIFF a block
# only becomes SKIP if holding the previous reconstruction costs less than
# 28,665 units of distortion, so a frame with nothing on screen to hold -- the
# first frame of a stream, or a scene cut -- stays fully coded and comes out at
# the all-V1 floor of 110.6% (FINDINGS 28.5). Such a frame is emitted LATE on
# purpose, exactly as a frame that will not fit at LAM_CLIFF is emitted over
# budget. Freezing a cut to make the deadline would be the worse failure.
MU_CLIFF = 100.0
MU_FLOOR = 1e-4                    # bisection is geometric, so lo must be > 0

# The hard per-frame decode budget. NOT a bucket: bytes can be banked in the
# player's ring buffer, but there is no double buffer to decode ahead into, so
# a frame that misses its deadline is simply late. FINDINGS 28.
FRAME_CYCLES = 10_000_000 / 12.0

AUDIO_KBPS = 7.8          # MSM6258 ADPCM 15.6kHz mono -- comes out of the same budget


def frame_budget(kbps, fps=12, audio=AUDIO_KBPS):
    """bytes per video frame after audio takes its cut"""
    return (kbps - audio) * 1024.0 / fps


def _search_lam(ctx, allow, lam_lo, lam_hi, iters=12, mu=0.0):
    """Smallest lam (=> best quality) whose frame fits `allow` bytes.

    Payload size is non-increasing in lam -- raising lam can only move a block
    to a mode that costs no more -- so bisection is sound. Geometric bisection,
    because lam spans three decades and the interesting range is multiplicative.

    Returns (lam, mode, size, overrun). `overrun` is True when even lam_hi does
    not fit: that frame is emitted over budget on purpose. Past the FINDINGS 15
    cliff a frame is not rate-controlled, it is destroyed, so a visible overrun
    is the better failure (FINDINGS 26.2)."""
    mode, sz = H.decide(ctx, lam_lo, mu)
    if sz <= allow:
        return lam_lo, mode, sz, False
    mode_hi, sz_hi = H.decide(ctx, lam_hi, mu)
    if sz_hi > allow:
        return lam_hi, mode_hi, sz_hi, True
    lo, hi = lam_lo, lam_hi                  # lo does not fit, hi does
    best = (lam_hi, mode_hi, sz_hi)
    for _ in range(iters):
        mid = float(np.sqrt(lo * hi))
        mode_m, sz_m = H.decide(ctx, mid, mu)
        if sz_m <= allow:
            hi = mid; best = (mid, mode_m, sz_m)
        else:
            lo = mid
    return best[0], best[1], best[2], False


def _search_mu(ctx, allow, lam_lo, lam_hi, cyc_budget, iters=10):
    """Smallest mu whose frame fits BOTH budgets: `allow` bytes and
    `cyc_budget` 68000 cycles.

    Two controllers, one nested inside the other, because the constraints are
    not separable. Raising mu moves blocks to cheaper-to-DECODE modes, which
    usually also shrinks the frame -- but not always: RAW is 400 cycles against
    V4's 448 and 16 bytes against 4, so mu can buy cycles by SPENDING bytes
    (FINDINGS 28.8). So every mu step re-runs the lam bisection and the byte
    budget is enforced at the mu that is actually chosen.

    Cost is scored with H.cycles(), the exact clustered rule, NOT with the
    per-block ranking constant the decision uses -- see vq_hybrid's note on
    SKIP. The controller therefore converges on what the 68000 will really do.

    Monotonicity: at a fixed lam, raising mu can only move a block to a mode
    that costs no more cycles, and it can only ADD to a SKIP cluster, so frame
    cycles are non-increasing in mu. The nested lam re-search can perturb that
    at the margin (a smaller frame permits a smaller lam, which buys quality
    back and can cost a few cycles), so the bisection keeps the best FEASIBLE
    point it has actually seen rather than trusting the invariant.

    Returns (mu, lam, mode, size, cyc, over_bytes, over_cycles)."""
    lam, mode, sz, ovr = _search_lam(ctx, allow, lam_lo, lam_hi, mu=0.0)
    cyc = H.cycles(mode)
    if cyc <= cyc_budget:
        return 0.0, lam, mode, sz, cyc, ovr, False

    lam_h, mode_h, sz_h, ovr_h = _search_lam(ctx, allow, lam_lo, lam_hi, mu=MU_CLIFF)
    cyc_h = H.cycles(mode_h)
    if cyc_h > cyc_budget:                   # cannot fit even frozen: emit late
        return MU_CLIFF, lam_h, mode_h, sz_h, cyc_h, ovr_h, True

    lo, hi = MU_FLOOR, MU_CLIFF              # lo overruns, hi fits
    best = (MU_CLIFF, lam_h, mode_h, sz_h, cyc_h, ovr_h)
    for _ in range(iters):
        mid = float(np.sqrt(lo * hi))
        lam_m, mode_m, sz_m, ovr_m = _search_lam(ctx, allow, lam_lo, lam_hi, mu=mid)
        cyc_m = H.cycles(mode_m)
        if cyc_m <= cyc_budget:
            hi = mid; best = (mid, lam_m, mode_m, sz_m, cyc_m, ovr_m)
        else:
            lo = mid
    return (*best, False)


def encode_rate_controlled(m, target_kbps, fps=12, bucket_frames=8,
                           lam_lo=1.0, lam_hi=LAM_CLIFF, prefill=0.0,
                           steps=None, verbose=False, cycle_budget=None):
    """Per-frame lam search under a leaky bucket, driving the encoder ONE FRAME
    AT A TIME and feeding back the frame actually emitted.

    That feedback is the whole point. The previous implementation encoded the
    sequence once per lam and then picked frames off the resulting ladder; the
    codec is temporally recursive, so frames picked from different rungs
    reference reconstructions the decoder never saw -- 111 of 120 frames drifted,
    worst frame 43.4% (FINDINGS 26.1). `tools/analysis/09_ratectl_drift.py` is
    the regression test and must report zero drifting frames.

    lam_lo is a QUALITY FLOOR, not a starting guess: rate control here only ever
    spends less than the fixed-lam profile, never more, so it cannot regress
    content that already fits. Pass lam_lo=1.0 to let quiet frames spend the
    whole allowance instead.

    `prefill` is how full the player's buffer is assumed to be when the scene
    starts, as a fraction of the bucket. 0.0 (the default) is the conservative
    assumption -- a cold buffer after a seek -- and is what FINDINGS 21 verified
    needs no prefill to avoid underflow. It costs a startup transient: the first
    `bucket_frames` frames cannot draw on a bank they have not accumulated yet,
    so a clip shorter than a few bucket depths lands UNDER target. That is an
    artefact of the clip length, not of the content; see FINDINGS 27.5.

    DO NOT raise `prefill` to make a target look met. It works by permitting an
    overshoot of cap/nframes: measured, prefill=1.0 takes the Singe window from
    109.5 to 116.3 KB/s against a 110 ceiling, and on a 14-frame clip it
    disables rate control entirely because the bucket is larger than the clip.

    `cycle_budget` adds the SECOND controller (session 8): a hard per-frame
    68000 decode ceiling, bisected on `mu` inside the lam search. None (the
    default) leaves it off and reproduces session 6 exactly, which is what
    keeps tools/analysis/09_ratectl_drift.py comparable. Pass
    FRAME_CYCLES for the 12fps stock-68000 budget.

    `steps` is accepted and ignored -- there is no ladder any more.
    """
    if steps is not None and verbose:
        print("  note: `steps` is ignored; lam is now bisected per frame")
    budget = frame_budget(target_kbps, fps)
    cap = bucket_frames * budget
    bucket = prefill * cap    # banked bytes; bounded by the player's buffer both ways
    out = dict(recon=[], modes=[], sizes=[], lam=[], l1=[], l4g=[], overrun=[],
               mu=[], cycles=[], late=[])
    prev = None
    for f in range(len(m["idx"])):
        ctx = H.frame_ctx(m, f, prev)
        allow = budget + bucket
        if cycle_budget is None:
            lam, mode, sz, ovr = _search_lam(ctx, allow, lam_lo, lam_hi)
            mu, cyc, late = 0.0, H.cycles(mode), False
        else:
            mu, lam, mode, sz, cyc, ovr, late = _search_mu(
                ctx, allow, lam_lo, lam_hi, cycle_budget)
        rec = H.paint(m, ctx, mode)
        bucket = float(np.clip(bucket + budget - sz, -cap, cap))
        out["recon"].append(rec); out["modes"].append(mode)
        out["sizes"].append(sz); out["lam"].append(lam); out["overrun"].append(ovr)
        out["mu"].append(mu); out["cycles"].append(cyc); out["late"].append(late)
        out["l1"].append(ctx["sym"]["l1"]); out["l4g"].append(ctx["sym"]["l4g"])
        prev = rec
        if verbose:
            print(f"  f{f:04d} lam={lam:8.2f} mu={mu:8.4f} {sz:7.0f} B "
                  f"(allow {allow:7.0f}) {100*cyc/FRAME_CYCLES:5.1f}% cpu"
                  f"{' OVER' if ovr else ''}{' LATE' if late else ''}")
    return dict(recon=out["recon"], modes=out["modes"],
                sizes=np.array(out["sizes"]), lam=np.array(out["lam"]),
                l1=out["l1"], l4g=out["l4g"], overrun=np.array(out["overrun"]),
                mu=np.array(out["mu"]), cycles=np.array(out["cycles"]),
                late=np.array(out["late"]),
                nb=m["nb"], budget=budget, cap=cap, cycle_budget=cycle_budget)


def summarise(m, enc, target_kbps, fps=12):
    import vq as VQ
    pal = m["pal"]
    rec = [pal[i] for i in enc["recon"]]
    src = [pal[i] for i in m["idx"]]
    p = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], rec)])
    pp = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], src)])
    sz = enc["sizes"]
    mo = np.concatenate(enc["modes"])
    d = dict(target=target_kbps, psnr=p, pal=pp, loss=pp - p,
             mean_B=sz.mean(), max_B=sz.max(), budget=enc.get("budget", 0.0),
             kbps=sz.mean() * fps / 1024 + AUDIO_KBPS,
             over=100.0 * np.mean(sz > enc.get("budget", np.inf)),
             skip=100 * (mo == 0).mean(), v1=100 * (mo == 1).mean(),
             v4=100 * (mo == 2).mean(), raw=100 * (mo == 3).mean())
    if "cycles" in enc:
        cy = np.asarray(enc["cycles"])
        d.update(cyc_med=float(np.median(cy)), cyc_max=float(cy.max()),
                 cyc_p90=float(np.percentile(cy, 90)),
                 cpu_miss=int((cy > FRAME_CYCLES).sum()),
                 mu_med=float(np.median(enc["mu"])),
                 mu_max=float(np.asarray(enc["mu"]).max()),
                 late=int(np.asarray(enc.get("late", [])).sum()))
    if "lam" in enc:
        lam = enc["lam"]
        d.update(lam_med=float(np.median(lam)), lam_max=float(lam.max()),
                 lam_p90=float(np.percentile(lam, 90)),
                 # a frame that could not fit even at the cliff: emitted over
                 # budget on purpose rather than destroyed
                 overrun=int(np.asarray(enc.get("overrun", [])).sum()))
    return d