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

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.
"""
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
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")
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()))

rec_b = P.write(a.out, W, H, a.fps, recs, palette_last=a.palette_last)
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}: DLXP1 {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")
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})")