entropy/ans.py

"""Chunk-local dictionary ANS coding.

Each chunk builds its own frequency histogram and uses it to drive zstd.
Scenes have very different token distributions, so adapting per chunk
improves compression significantly.

Tokens are 18-bit packed before zstd to avoid waste.
"""

import struct
import zlib
from collections import Counter
from mediatok.entropy.pack import pack_tokens, unpack_tokens


def ans_encode(tokens: list[int]) -> bytes:
    """Encode tokens using chunk-local frequency table + 18-bit pack + zstd."""
    if not tokens:
        return b""

    freq = Counter(tokens)

    meta = struct.pack("<II", len(tokens), len(freq))
    for sym in sorted(freq.keys()):
        meta += struct.pack("<II", sym, freq[sym])

    packed = bytes(pack_tokens(tokens, 18))
    compressed = zlib.compress(packed, level=3)

    return struct.pack("<I", len(meta)) + meta + compressed


def ans_decode(payload: bytes) -> list[int]:
    if not payload:
        return []

    hlen = struct.unpack_from("<I", payload, 0)[0]
    meta = payload[4 : 4 + hlen]
    num_syms, n_sym = struct.unpack_from("<II", meta, 0)
    if num_syms == 0:
        return []
    if n_sym == 0:
        return []

    compressed = payload[4 + hlen :]
    packed = zlib.decompress(compressed)
    return unpack_tokens(list(packed), num_syms, 18)