#!/usr/bin/env python3
"""Encode one scene to the PACKED container -- ROADMAP K2.

    python3 tools/encoder/pack.py <frames_dir> <out.dlxp> [--fps 12]
                                  [--nframes N] [--palette-last] [--no-palette]
                                  [--scene-palette] [--audio tmp/au_singe.raw]

WHAT IS NOT HERE IS THE POINT.  No VQ, no codebooks, no mode map, no rate
control, no `lam`, no leaky bucket, no span geometry -- `encode.py` is 452 lines
and 95% of its wall clock is k-means.  A packed frame is a palette and a
picture, and both are geometry.  FINDINGS 61: at the 9 clk/B dual-address floor
the codec is 110.4% of a 12 fps frame and this is 55.2%, so the thing that
replaces the codec is also the thing that is simpler than it.

THERE IS NO RATE CONTROL BECAUSE THERE IS NO RATE LEVER.  A codec's bitrate is
adjustable; a literal frame's is geometry.  The wire cost of this container is
fixed by W, H and fps and nothing an encoder does can move it, which is exactly
why FINDINGS 61.6 says the medium question decides which player exists.  A
`--kbps` argument here would be a lie of the shape FINDINGS 50 removed from the
rest of the tree.

THE QUANTISER IS THIS PROJECT'S, NOT PIL'S DEFAULT PATH.  61.9's +4.89 dB was
measured with `18_text_plane_16col.py`'s free 256-colour MEDIANCUT and was filed
as "a direction, not the player's number" (risk 2 in the session 30 handoff).
`vq.frame_palette` is what actually ships it: 254 colours, index 0 held free for
the transparency key, black at 255.  `tools/analysis/30_packed_container.py`
re-derives the figure against this and charges the GRB555 word on top.

AND `--audio` MAKES IT A DLXP2, WHICH IS THE ONE PLACE THE FOUR ADPCM AXES ARE
CHOSEN.  It encodes for `adpcm.CHIP` -- the datasheet's per-term delta, the LOW
nibble of a byte first, a 10-bit clamp, the accumulator at -2 -- because that is
the set FINDINGS 66 measured out of the machine's own chip through the machine's
own DMA channel, and encoding for any other set costs up to 25.7 dB.  It does
NOT use adpcm.py's module defaults, which are ffmpeg's on purpose so that
tools/bench/verify_adpcm.py stays a check against an independent implementation.
The axes go in the header, so a player never has to be told.
"""
import argparse, glob, os, sys
import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                "..", "bench"))
import vq as VQ
import dlxp as P
import adpcm
from dlxload import pack_palette

ap = argparse.ArgumentParser()
ap.add_argument("frames_dir")
ap.add_argument("out")
ap.add_argument("--fps", type=int, default=12)
ap.add_argument("--nframes", type=int, default=None,
                help="encode only the first N frames (the gate window is 120)")
ap.add_argument("--palette-last", action="store_true",
                help="put the palette after the picture in every record. "
                     "FINDINGS 62.5 -- a design choice, not a default to inherit")
ap.add_argument("--no-palette", action="store_true",
                help="picture only, 49,152 B a record. NOT a shipping option: "
                     "it re-imposes the scene palette the codec is capped by")
ap.add_argument("--scene-palette", action="store_true",
                help="one palette for the whole scene, repeated in every record "
                     "-- the CONTROL for 61.9's per-frame claim")
ap.add_argument("--audio", default=None,
                help="raw s16le mono at --audio-hz, from extract_audio.py. "
                     "Makes the output a DLXP2: the stream is encoded for "
                     "adpcm.CHIP and interleaved on the 65.3 cadence")
ap.add_argument("--audio-hz", type=int, default=15625,
                help="the CHIP's rate, and the rate the .raw was resampled to")
ap.add_argument("--cadence", type=int, default=P.CADENCE_F,
                help="frames between audio lumps (65.3's sweep picks 11)")
a = ap.parse_args()

files = sorted(glob.glob(f"{a.frames_dir}/f*.png"))
if a.nframes:
    files = files[:a.nframes]
if not files:
    sys.exit(f"no f*.png in {a.frames_dir}")
rgb = [np.asarray(Image.open(f).convert("RGB")) for f in files]
H, W = rgb[0].shape[:2]
if any(r.shape[:2] != (H, W) for r in rgb):
    sys.exit("frames are not all the same size")

# The scene-palette control shares ONE palette across every record, which is the
# constraint the codec cannot escape (61.9) and this format merely chooses not
# to inherit.  It is built by the same routine the codec uses, with black at 0
# moved to 255 and index 0 vacated, so the only variable between the two runs is
# per-frame versus scene-wide.
scene = None
if a.scene_palette:
    # 255, not 256: reserve_black spends one entry on black and the packed
    # layout needs the OTHER end free too, so the picture gets 254 either way.
    ref, spal = VQ.scene_palette(rgb, colors=255, reserve_black=True)
    sidx = VQ.palettise(rgb, ref)
    # 0 -> 255: the packed layout needs index 0 free and black displayed at 255.
    spal = np.vstack([np.zeros((1, 3), np.uint8), spal[1:],
                      np.zeros((1, 3), np.uint8)])
    scene = (spal, [np.where(i == 0, np.uint8(255), i) for i in sidx])

recs, psnr_pal = [], []
for n, src in enumerate(rgb):
    if scene is not None:
        pal, idx = scene[0], scene[1][n]
    else:
        pal, idx = VQ.frame_palette(src)
    if (idx == 0).any():
        sys.exit(f"frame {n}: index 0 is the transparency key and got used")
    palw = None
    if not a.no_palette:
        palb, _dark, rendered = pack_palette(pal)
        palw = palb.tobytes()
    psnr_pal.append(VQ.psnr(src, pal[idx]))
    recs.append((palw, P.pack_picture(idx).tobytes()))

# THE AUDIO, AND IT IS ENCODED FOR THE CHIP AND NOT FOR ffmpeg.  The source is
# s16le because that is what ffmpeg resamples to; the chip's word is 12 bits
# signed, so the shift is a requantisation and not a format conversion, and it
# is the same one tools/bench/verify_adpcm.py makes.
audio = None
if a.audio:
    import struct as _struct
    raw = open(a.audio, "rb").read()
    pcm = _struct.unpack("<%dh" % (len(raw) // 2), raw)
    src12 = [max(-2048, min(2047, x >> 4)) for x in pcm]
    need = len(files) * a.audio_hz // (2 * a.fps)
    nib = adpcm.encode(src12, variant=adpcm.CHIP["variant"],
                       init=adpcm.CHIP["init"], bits=adpcm.CHIP["bits"])
    data = adpcm.pack(nib, order=adpcm.CHIP["order"])[:need]
    # What the encoder thinks it wrote, checked by DECODING it with the same
    # four axes -- the encoder runs its decoder inside its own loop, so this is
    # not circular in the way it looks: it is the PACKED bytes going back
    # through unpack(), which is where a nibble-order slip would land.
    back = adpcm.decode(adpcm.unpack(data, len(src12), order=adpcm.CHIP["order"]),
                        variant=adpcm.CHIP["variant"], init=adpcm.CHIP["init"],
                        bits=adpcm.CHIP["bits"])
    ref = src12[:len(back)]
    e = [(x - y) ** 2 for x, y in zip(ref, back)]
    sig = sum(x * x for x in ref)
    snr = 10 * np.log10(sig / sum(e)) if sum(e) else float("inf")
    audio = dict(data=data, hz=a.audio_hz, F=a.cadence, **adpcm.CHIP)

rec_b = P.write(a.out, W, H, a.fps, recs, palette_last=a.palette_last,
                audio=audio)
d = P.DLXP(a.out)                        # re-read: every invariant is checked
if d.nframes != len(recs):
    sys.exit("writer and reader disagree about the frame count")

kind = ("scene palette" if a.scene_palette else "per-frame palette")
if a.no_palette:
    kind += ", NONE in the record"
print(f"{a.out}: DLXP{P.VERSION} {W}x{H} {a.fps}fps {d.nframes} frames, {kind}"
      f"{', palette LAST' if a.palette_last else ''}")
print(f"  record {rec_b:,} B = {rec_b // P.SECTOR} sectors exactly, "
      f"file {os.path.getsize(a.out):,} B")
if d.has_audio:
    print(f"  DLXP2: audio {d.aud_bytes:,} B at {d.aud_hz:,} Hz, SNR {snr:.2f} dB, "
          f"cadence F={d.cad_f} A={d.cad_a} ({d.n_lumps} lumps)")
    print(f"  the four axes, in the header: "
          + ", ".join(f"{k}={v}" for k, v in d.decoder().items()))
    # STEADY STATE, not the file: the last lump of a 120-frame window feeds 4
    # frames out of 11 and occupies 14 sectors either way, so a whole-file
    # padding figure is a boundary effect of the WINDOW and would change with
    # its length.  65.3's 0.09% is the cadence's, and the cadence is the thing.
    grp = d.cad_f * d.aud_hz / (2 * d.fps)
    print(f"  lump payload {P.lump_bytes(0, d.cad_f, d.fps, d.aud_hz):,}.."
          f"{P.lump_bytes(2, d.cad_f, d.fps, d.aud_hz):,} B of {d.cad_a*P.SECTOR:,} "
          f"-- {100*(d.cad_a*P.SECTOR-grp)/grp:.3f}% padding steady state, and the "
          f"payload is NOT the lump (FINDINGS 67)")
    print(f"  wire {d.video_kbps():.1f} + {d.audio_kbps():.2f} = {d.kbps():.1f} KB/s "
          f"-- FIXED by geometry, there is no lever")
else:
    print(f"  wire {d.kbps():.1f} KB/s -- FIXED by geometry, there is no lever")
print(f"  palette-domain PSNR vs the 24-bit source: "
      f"{np.mean(psnr_pal):.2f} dB (min {np.min(psnr_pal):.2f})")