"""
Real Witten-Neal-Cleary arithmetic coder, order-2 ADAPTIVE byte model.
Round-trips a slice of enwik8, ASSERTS lossless, reports actual bits/byte.
Purpose: prove the ~3.09 bits/byte order-2 estimate is achievable by a genuine
lossless coder (not just a theoretical entropy), on the same trap-free footing
as the char-poly audit. Slice keeps pure-Python runtime sane; rate is the point.
"""
import sys, numpy as np

PATH = sys.argv[1] if len(sys.argv) > 1 else "enwik8"
SLICE = int(sys.argv[2]) if len(sys.argv) > 2 else 3_000_000

with open(PATH, "rb") as f:
    data = f.read(SLICE)
n = len(data)

# ---- WNC 32-bit arithmetic coder ----
BITS = 32
TOP = (1 << BITS) - 1
HALF = 1 << (BITS - 1)
QTR = 1 << (BITS - 2)
TQTR = 3 << (BITS - 2)
MAX_TOT = 1 << 16          # keep total freq below 2^16 << 2^30 precision bound

NCTX = 1 << 16            # order-2 contexts (prev two bytes)
# adaptive freqs, init 1 (Laplace). uint32 so a row can't overflow before rescale.
freq = np.ones((NCTX, 256), dtype=np.uint32)

def rescale(row):
    row >>= 1
    row |= 1  # keep every symbol >=1 so nothing becomes uncodeable

# ---------- ENCODE ----------
low, high = 0, TOP
pending = 0
out = bytearray()
bitbuf = 0
nbits = 0
def put(bit):
    global bitbuf, nbits
    bitbuf = (bitbuf << 1) | bit
    nbits += 1
    if nbits == 8:
        out.append(bitbuf)
        bitbuf = 0
        nbits = 0
def emit(bit):
    global pending
    put(bit)
    while pending:
        put(bit ^ 1)
        pending -= 1

ctx = 0
for sym in data:
    row = freq[ctx]
    tot = int(row.sum())
    cum = int(row[:sym].sum())
    f = int(row[sym])
    rng = high - low + 1
    high = low + (rng * (cum + f)) // tot - 1
    low = low + (rng * cum) // tot
    while True:
        if high < HALF:
            emit(0)
        elif low >= HALF:
            emit(1); low -= HALF; high -= HALF
        elif low >= QTR and high < TQTR:
            pending += 1; low -= QTR; high -= QTR
        else:
            break
        low <<= 1
        high = (high << 1) | 1
    # adapt
    row[sym] += 32
    if tot + 32 >= MAX_TOT:
        rescale(row)
    ctx = ((ctx << 8) | sym) & 0xFFFF
# flush
pending += 1
emit(1 if low >= QTR else 0)
if nbits:
    out.append(bitbuf << (8 - nbits))
enc = bytes(out)
enc_bits = len(enc) * 8

# ---------- DECODE ----------
freq = np.ones((NCTX, 256), dtype=np.uint32)
low, high = 0, TOP
bit_iter_pos = 0
total_bits = len(enc) * 8
def getbit():
    global bit_iter_pos
    if bit_iter_pos < total_bits:
        byte = enc[bit_iter_pos >> 3]
        b = (byte >> (7 - (bit_iter_pos & 7))) & 1
    else:
        b = 0
    bit_iter_pos += 1
    return b

code = 0
for _ in range(BITS):
    code = (code << 1) | getbit()

dec = bytearray()
ctx = 0
for _ in range(n):
    row = freq[ctx]
    tot = int(row.sum())
    rng = high - low + 1
    # find symbol whose cumulative interval contains 'code'
    val = ((code - low + 1) * tot - 1) // rng
    cs = np.cumsum(row)
    sym = int(np.searchsorted(cs, val, side='right'))
    cum = int(cs[sym - 1]) if sym > 0 else 0
    f = int(row[sym])
    high = low + (rng * (cum + f)) // tot - 1
    low = low + (rng * cum) // tot
    while True:
        if high < HALF:
            pass
        elif low >= HALF:
            low -= HALF; high -= HALF; code -= HALF
        elif low >= QTR and high < TQTR:
            low -= QTR; high -= QTR; code -= QTR
        else:
            break
        low <<= 1
        high = (high << 1) | 1
        code = (code << 1) | getbit()
    dec.append(sym)
    row[sym] += 32
    if tot + 32 >= MAX_TOT:
        rescale(row)
    ctx = ((ctx << 8) | sym) & 0xFFFF

ok = bytes(dec) == data
print(f"slice bytes        : {n}")
print(f"encoded bytes      : {len(enc)}")
print(f"actual rate        : {enc_bits / n:6.3f} bits/byte")
print(f"lossless round-trip: {'PASS' if ok else 'FAIL'}")
if not ok:
    for i in range(n):
        if dec[i] != data[i]:
            print(f"  first mismatch at {i}: {dec[i]} != {data[i]}")
            break
print("DONE_AC")