#!/usr/bin/env python3
"""Lay a DLX1 container out the way src/player/decode.s expects to find it.

    python3 tools/bench/prep_dlx.py <in.dlx> [--out tmp/decode]

Writes <out>_data.bin (one blob Lua pushes into emulated RAM) and <out>_meta.lua
(sizes, per-frame record offsets, and the timing anchors).

Two things happen here that the shipping player would do at load time on the
68000 itself, and are therefore NOT part of the per-frame cost being measured:

  * codebook expansion to word-per-pixel form.  CB1 -> 32 B/entry, CB4 -> 8 B,
    so the inner loop scales an index with lsl.w #5 / #3 and movems the result
    straight into GVRAM with no unpacking.  8 KB + 2 KB of the 2 MB.
  * palette packing to GGGGGRRRRRBBBBBI with the shared LSB I chosen PER ENTRY
    by minimum squared error (FINDINGS 23.3, worth 1.96 dB).

The encoder still emits 24-bit palettes and does not reserve a black entry
(known gap, docs/STATUS.md), so the letterbox here is filled with whatever
palette entry is closest to black rather than a true reserved black.  That is
cosmetic and outside the active 256x192 area the decoder is judged on.

A synthetic all-SKIP frame is appended to the stream.  No real frame is all
SKIP, but it prices the mode-header walk on its own -- the per-block cost the
"76.6% x non-SKIP fraction" model in FINDINGS 24.5 leaves out entirely.
"""
import sys, os, argparse
sys.path.insert(0, "tools/encoder")
import numpy as np
from dlx import DLX

# The harness loads the WHOLE container into emulated RAM at STREAM=0x30000 and
# the target is a stock 2 MB machine, so there is a hard ceiling on how much of
# a stream can be verified in one pass. The shipping player streams from disk
# into a ring buffer and has no such limit; this is a property of the test rig.
# A `scsi` window overruns it -- 2.84 MB of stream ends at 0x2E591C, 940 KB past
# the 0x200000 top of RAM -- so the frame list is truncated to what fits and the
# truncation is announced. Verifying a prefix is still a real test: SKIP blocks
# make every frame a claim about the one before it.
STREAM_BASE = 0x30000
RAM_TOP     = 0x200000
MARGIN      = 0x8000          # stack, flags, codebooks live below STREAM_BASE

ap = argparse.ArgumentParser()
ap.add_argument("container")
ap.add_argument("--out", default="tmp/decode")
ap.add_argument("--ram", type=lambda v: int(v, 0), default=RAM_TOP,
                help="top of emulated RAM (default 0x200000, a stock 2 MB machine)")
ap.add_argument("--all-frames", action="store_true",
                help="do NOT truncate to what fits in RAM (the loader will write "
                     "past the top of memory and the decoder will read garbage)")
a = ap.parse_args()

d = DLX(a.container)
if d.idx_bytes != 1:
    sys.exit("2-byte codebook indices: decode.s assumes 1 (k<=256)")

# --- codebooks, expanded to one WORD per pixel (high byte is discarded by
#     gvram_w, so it is left zero and never has to be cleared)
cb1 = np.zeros((d.k1, 16, 2), np.uint8); cb1[:, :, 1] = d.cb1.reshape(d.k1, 16)
cb4 = np.zeros((d.k4, 4, 2), np.uint8);  cb4[:, :, 1] = d.cb4.reshape(d.k4, 4)

# --- palette words, I chosen per entry (identical maths to verify_frame256.py)
pal = d.pal.astype(int)
p6 = lambda v: ((v << 2) | (v >> 4)) & 0xFF
f = pal >> 3
render = lambda I: p6((f << 1) | I[:, None])
I = (((render(np.ones(256, int)) - pal) ** 2).sum(1)
     < ((render(np.zeros(256, int)) - pal) ** 2).sum(1)).astype(int)
words = (f[:, 1] << 11) | (f[:, 0] << 6) | (f[:, 2] << 1) | I
palb = np.zeros((256, 2), np.uint8)
palb[:, 0], palb[:, 1] = words >> 8, words & 0xFF
dark = int(((render(I).astype(int)) ** 2).sum(1).argmin())

# --- frame stream: [u32 len][modes][payload] per frame, each record start
#     rounded up to a 4-byte boundary.
#
#     This padding is not cosmetic.  Payload lengths are arbitrary, so laid end
#     to end the records land on odd addresses, and `move.l (a0)+` at an odd
#     address is an ADDRESS ERROR on a 68000 -- it vectors into the IPL rather
#     than reading slowly.  The container as written by encode.py is unaligned,
#     so this loader realigns it; the encoder should carry the padding itself
#     (FINDINGS 28.3).  It costs at most 3 bytes per frame -- 36 B/s at 12fps,
#     against a 110 KB/s budget.
budget = a.ram - STREAM_BASE - MARGIN
stream, rec_off, pad = bytearray(), [], 0
dropped = 0
for (o, n) in d.frames:
    while len(stream) % 4:
        stream += b"\0"; pad += 1
    if not a.all_frames and len(stream) + 4 + n > budget:
        dropped = d.nframes - len(rec_off)
        break
    rec_off.append(len(stream))
    stream += n.to_bytes(4, "big") + d.raw[o:o + n]
NFRAMES = len(rec_off)
if dropped:
    print(f"  TRUNCATED: {NFRAMES}/{d.nframes} frames fit in RAM "
          f"(stream budget {budget:,} B at 0x{STREAM_BASE:X} under a "
          f"{a.ram/1024/1024:.0f} MB machine); {dropped} frames dropped.\n"
          f"  This is the TEST RIG's limit, not the player's -- the player "
          f"streams into a ring buffer.")

# Synthetic single-mode frames.  No real frame is all one mode, but the mix is
# exactly what the "76.6% x non-SKIP fraction" model of FINDINGS 24.5 assumes
# away: it prices every non-SKIP block as one V1-style burst.  These four price
# the modes separately, which is the only way to see which one is expensive.
synth = {}
for name, mo, per in (("all-SKIP", 0, 0), ("all-V1", 1, 1),
                      ("all-V4", 2, 4), ("all-RAW", 3, 16)):
    while len(stream) % 4:
        stream += b"\0"; pad += 1
    synth[name] = len(stream)
    hdr = bytes([mo * 0x55] * d.mode_bytes)
    stream += (d.mode_bytes + d.nb * per).to_bytes(4, "big") + hdr + bytes(d.nb * per)

# --- timing anchors: the distribution, not its mean (FINDINGS 25.6's lesson)
ns = np.array([100 * (d.modes(i) != 0).mean() for i in range(NFRAMES)])
order = np.argsort(ns)
pick = {
    "min non-SKIP  %.1f%%" % ns[order[0]]: int(order[0]),
    "median        %.1f%%" % np.median(ns): int(order[len(order)//2]),
    "p90           %.1f%%" % ns[order[int(.9*len(order))]]: int(order[int(.9*len(order))]),
    "max non-SKIP  %.1f%%" % ns[order[-1]]: int(order[-1]),
}
anchors = [(n, rec_off[i], float(ns[i])) for n, i in pick.items()]
for name in ("all-SKIP", "all-V1", "all-V4", "all-RAW"):
    anchors.append((f"synthetic {name}", synth[name],
                    0.0 if name == "all-SKIP" else 100.0))

blob = cb1.tobytes() + cb4.tobytes() + palb.tobytes() + bytes(stream)
open(a.out + "_data.bin", "wb").write(blob)

with open(a.out + "_meta.lua", "w") as fh:
    fh.write("-- generated by tools/bench/prep_dlx.py -- do not edit\nreturn {\n")
    fh.write(f"  W={d.W}, H={d.H}, fps={d.fps}, nframes={NFRAMES},\n")
    fh.write(f"  k1={d.k1}, k4={d.k4}, dark={dark},\n")
    fh.write(f"  cb1_len={cb1.nbytes}, cb4_len={cb4.nbytes}, pal_len={palb.nbytes},\n")
    fh.write(f"  stream_len={len(stream)},\n")
    fh.write("  anchors={\n")
    for n, o, frac in anchors:
        fh.write(f'    {{name="{n}", off={o}, frac={frac:.1f}}},\n')
    fh.write("  },\n}\n")

print(f"{a.container}: {d.nframes} frames, {d.W}x{d.H}, k1={d.k1} k4={d.k4}")
print(f"  cb1 {cb1.nbytes} B + cb4 {cb4.nbytes} B expanded, palette {palb.nbytes} B, "
      f"stream {len(stream)} B  ->  {a.out}_data.bin ({len(blob)} B)")
print(f"  non-SKIP blocks/frame: median {np.median(ns):.1f}%  "
      f"p90 {np.percentile(ns,90):.1f}%  max {ns.max():.1f}%")
print(f"  darkest palette entry: index {dark} -> {tuple(render(I)[dark])}")
# A DLX2 container already carries this padding (FINDINGS 28.3 closed, session
# 9), so the realignment above re-derives bytes that were already there and the
# loader is doing no work. On a DLX1 container it is load-bearing: 94 of 120
# record starts land on odd addresses, and each one is an address error.
src_bad = sum(1 for (o, _) in d.frames[:NFRAMES] if (o - 4) % 4)
print(f"  4-byte record alignment cost {pad} B over {NFRAMES} frames "
      f"({pad / NFRAMES:.2f} B/frame = {pad / NFRAMES * d.fps:.0f} B/s)")
print(f"  source container is DLX{d.version}: {src_bad}/{NFRAMES} record starts "
      f"unaligned" + ("  -- this loader is what makes it decodable"
                      if src_bad else "  -- the container carries its own padding"))