#!/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

ap = argparse.ArgumentParser()
ap.add_argument("container")
ap.add_argument("--out", default="tmp/decode")
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.
stream, rec_off, pad = bytearray(), [], 0
for (o, n) in d.frames:
    while len(stream) % 4:
        stream += b"\0"; pad += 1
    rec_off.append(len(stream))
    stream += n.to_bytes(4, "big") + d.raw[o:o + n]

# 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(d.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={d.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])}")
print(f"  4-byte record alignment cost {pad} B over {d.nframes} frames "
      f"({pad / d.nframes:.2f} B/frame = {pad / d.nframes * d.fps:.0f} B/s)")