#!/usr/bin/env python3
"""Encode one scene to the DLX bitstream, at a chosen quality profile.

    python3 tools/encoder/encode.py <frames_dir> <out.dlx> [--profile sasi|scsi]
                                    [--lam N] [--fps 12] [--preview out.png]

Container (little-endian is WRONG here -- the 68000 is big-endian, so every
multi-byte field is big-endian and the decoder can read it with a plain move.w):

  header, 32 bytes
    0   'DLX1'                 magic
    4   u16 width, u16 height
    8   u16 fps,   u16 nframes
    12  u16 k1,    u16 k4      codebook sizes
    16  u32 palette offset     (256 * 3 bytes, RGB888 -- the player converts
                                to the X68000's GRB555 at load time)
    20  u32 cb1 offset         (k1 * 16 bytes of palette indices)
    24  u32 cb4 offset         (k4 *  4 bytes)
    28  u32 frames offset
  then, per frame:
    u32 payload length, then
    ceil(nblocks*2/8) bytes of 2-bit mode headers, MSB-first, block raster order
    then payloads in block order: V1 -> 1 byte, V4 -> 4 bytes, RAW -> 16 bytes

Codebooks are emitted as palette INDICES, not pixels. The player expands them
once at load time into word-per-pixel form so the blitter can movem them
straight into GVRAM -- k1=1024 costs 1024*16*2 = 32 KB of the 2 MB.
"""
import argparse, struct, sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
import vq as VQ, vq_hybrid as H, ratectl as RC


def pack_modes(mode):
    """2 bits per block, MSB-first -- cheap for the 68000 to shift out."""
    n = len(mode)
    out = bytearray((n * 2 + 7) // 8)
    for i, m in enumerate(mode):
        out[i // 4] |= (int(m) & 3) << (6 - 2 * (i % 4))
    return bytes(out)


def frame_payload(mode, l1, l4g, src_idx, nbx):
    body = bytearray()
    for b, mo in enumerate(mode):
        if mo == 1:
            body += _idx(l1[b])
        elif mo == 2:
            for j in range(4):
                body += _idx(l4g[b][j])
        elif mo == 3:
            by, bx = divmod(b, nbx)
            body += src_idx[by*4:by*4+4, bx*4:bx*4+4].tobytes()
    return bytes(body)


def _idx(v):
    """codebook index: 1 byte if it fits, else big-endian u16.
    k>256 means 2-byte indices -- decided once by the header, not per block."""
    v = int(v)
    return bytes([v]) if _IDX_BYTES == 1 else struct.pack(">H", v)


_IDX_BYTES = 1


def main():
    global _IDX_BYTES
    ap = argparse.ArgumentParser()
    ap.add_argument("frames_dir"); ap.add_argument("out")
    ap.add_argument("--profile", choices=list(RC.PROFILES), default="sasi")
    ap.add_argument("--lam", type=float, default=None)
    ap.add_argument("--fps", type=int, default=12)
    ap.add_argument("--iters", type=int, default=16)
    ap.add_argument("--preview")
    a = ap.parse_args()

    prof = RC.PROFILES[a.profile]
    lam = a.lam if a.lam is not None else prof["lam"]
    k1, k4 = prof["k1"], prof["k4"]
    _IDX_BYTES = 1 if max(k1, k4) <= 256 else 2

    print(f"profile {a.profile}: {prof['desc']}")
    print(f"  target {prof['kbps']} KB/s, lam={lam}, k1={k1} k4={k4}, "
          f"{_IDX_BYTES}-byte indices")

    m = H.build(a.frames_dir, k1=k1, k4=k4, iters=a.iters)
    enc = H.encode(m, lam=lam)
    r = H.evaluate(m, enc, fps=a.fps)

    H_, W_ = m["H"], m["W"]; nbx = W_ // 4
    pal, idx = m["pal"], m["idx"]

    # re-derive the per-frame symbols the same way encode() did
    frames = []
    for f, im in enumerate(idx):
        B1 = H.blocks_of(im, pal, 4, 4); l1 = VQ.assign(B1, m["C1s"])
        B4 = H.blocks_of(im, pal, 2, 2); l4 = VQ.assign(B4, m["C4s"])
        q = H._group_2x2_into_4x4(np.arange(len(l4)), W_)
        l4g = l4[q].reshape(-1, 4)
        mode = enc["modes"][f]
        frames.append(pack_modes(mode) + frame_payload(mode, l1, l4g, im, nbx))

    palette = m["pal"][:256]
    if len(palette) < 256:
        palette = np.vstack([palette, np.zeros((256 - len(palette), 3), np.uint8)])
    pal_b = palette.astype(np.uint8).tobytes()
    cb1_b = m["cb1"].astype(np.uint8).tobytes()
    cb4_b = m["cb4"].astype(np.uint8).tobytes()

    off_pal = 32
    off_cb1 = off_pal + len(pal_b)
    off_cb4 = off_cb1 + len(cb1_b)
    off_frm = off_cb4 + len(cb4_b)
    hdr = (b"DLX1" + struct.pack(">HHHHHH", W_, H_, a.fps, len(idx), k1, k4)
           + struct.pack(">IIII", off_pal, off_cb1, off_cb4, off_frm))
    assert len(hdr) == 32, len(hdr)

    with open(a.out, "wb") as fh:
        fh.write(hdr); fh.write(pal_b); fh.write(cb1_b); fh.write(cb4_b)
        for p in frames:
            fh.write(struct.pack(">I", len(p))); fh.write(p)

    total = os.path.getsize(a.out)
    vid = sum(len(p) + 4 for p in frames)
    print(f"  wrote {a.out}: {total} B  "
          f"(header+tables {total-vid} B, video {vid} B)")
    print(f"  {vid/len(idx):.0f} B/frame -> {vid/len(idx)*a.fps/1024:.1f} KB/s video"
          f" + {RC.AUDIO_KBPS} KB/s audio = {vid/len(idx)*a.fps/1024+RC.AUDIO_KBPS:.1f} KB/s")
    print(f"  PSNR {r['psnr']:.2f} dB   palette ceiling {r['pal']:.2f} dB   "
          f"loss {r['loss']:.2f} dB")
    print(f"  modes: SKIP {r['skip']:.1f}%  V1 {r['v1']:.1f}%  "
          f"V4 {r['v4']:.1f}%  RAW {r['raw']:.1f}%")

    if a.preview:
        from PIL import Image
        f = len(idx) // 2
        gap = np.full((H_ * 3, 4, 3), 40, np.uint8)
        st = np.concatenate([VQ.zoom(m["rgb"][f], 3), gap,
                             VQ.zoom(pal[idx[f]], 3), gap,
                             VQ.zoom(pal[enc["recon"][f]], 3)], axis=1)
        Image.fromarray(st).save(a.preview)
        print(f"  preview -> {a.preview}  (source | palette ceiling | decoded)")


if __name__ == "__main__":
    main()