"""
Frozen-model + real arithmetic coder = lossless compressor, with HONEST
two-column accounting (the "weird machine" conservation law).

A decompressor is a program (Turing machine) that outputs the data. Its
self-contained size has two columns:
    TOTAL = |program: frozen model + coder|  +  |tape: arithmetic-coded residual|
Turing-completeness does not help: for the SAME data, making the model bigger
(higher order) shrinks the tape but grows the program. You can shuffle bits
between the columns; you cannot push the SUM below the data's entropy / K.

This script:
  - builds a FROZEN order-k byte model on a TRAIN split (no adaptation at code time)
  - compresses a held-out TEST split with a genuine Witten-Neal-Cleary coder
  - ASSERTS lossless round-trip
  - reports, per k: test bits/char, shipped model size (xz of the frozen table),
    tape bytes, self-contained TOTAL, and the amortized (model-shared) number.
  - baselines: xz / gzip on the same test bytes.
"""
import os, sys, pickle, lzma, gzip, numpy as np

def load_corpus(root, cap):
    buf = bytearray()
    for dp, _, fs in os.walk(root):
        for fn in fs:
            if fn.endswith((".md", ".txt", ".lean", ".py")):
                try:
                    with open(os.path.join(dp, fn), "rb") as f:
                        buf += f.read()
                except Exception:
                    pass
                if len(buf) >= cap:
                    return bytes(buf[:cap])
    return bytes(buf)

root = sys.argv[1] if len(sys.argv) > 1 else "/home/allaun/SilverSight"
data = load_corpus(root, 1_800_000)
split = len(data) - 150_000
train, test = data[:split], data[split:]
print(f"train {len(train)} bytes, test {len(test)} bytes (held out)\n")

# ---- WNC 32-bit arithmetic coder over a supplied freq vector ----
BITS = 32; TOP = (1 << BITS) - 1; HALF = 1 << 31; QTR = 1 << 30; TQ = 3 << 30

def freq_for(row_counts):
    """Deterministic positive freq vector from frozen counts (both sides identical)."""
    s = row_counts.sum()
    if s > 0:
        scaled = (row_counts * 4096.0 / s).astype(np.int64)
    else:
        scaled = np.zeros(256, dtype=np.int64)
    return scaled + 1  # Laplace: every symbol codeable; total < 4096+256 << 2^16

def encode(test, model, k):
    low, high, pending = 0, TOP, 0
    out = bytearray(); bitbuf = 0; nb = 0
    def put(b):
        nonlocal bitbuf, nb
        bitbuf = (bitbuf << 1) | b; nb += 1
        if nb == 8: out.append(bitbuf); bitbuf = 0; nb = 0
    def emit(b):
        nonlocal pending
        put(b)
        while pending: put(b ^ 1); pending -= 1
    ctx = 0; mask = (1 << (8 * k)) - 1 if k else 0
    for sym in test:
        row = model.get(ctx, ZERO)
        f = freq_for(row); tot = int(f.sum())
        cum = int(f[:sym].sum()); fs = int(f[sym])
        rng = high - low + 1
        high = low + rng * (cum + fs) // 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 < TQ: pending += 1; low -= QTR; high -= QTR
            else: break
            low <<= 1; high = (high << 1) | 1
        ctx = ((ctx << 8) | sym) & mask if k else 0
    pending += 1; emit(1 if low >= QTR else 0)
    if nb: out.append(bitbuf << (8 - nb))
    return bytes(out)

def decode(enc, model, k, n):
    low, high = 0, TOP; pos = 0; total = len(enc) * 8
    def getbit():
        nonlocal pos
        b = (enc[pos >> 3] >> (7 - (pos & 7))) & 1 if pos < total else 0
        pos += 1; return b
    code = 0
    for _ in range(BITS): code = (code << 1) | getbit()
    out = bytearray(); ctx = 0; mask = (1 << (8 * k)) - 1 if k else 0
    for _ in range(n):
        row = model.get(ctx, ZERO)
        f = freq_for(row); tot = int(f.sum())
        rng = high - low + 1
        val = ((code - low + 1) * tot - 1) // rng
        cs = np.cumsum(f); sym = int(np.searchsorted(cs, val, side='right'))
        cum = int(cs[sym - 1]) if sym > 0 else 0; fs = int(f[sym])
        high = low + rng * (cum + fs) // 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 < TQ: low -= QTR; high -= QTR; code -= QTR
            else: break
            low <<= 1; high = (high << 1) | 1; code = (code << 1) | getbit()
        out.append(sym); ctx = ((ctx << 8) | sym) & mask if k else 0
    return bytes(out)

ZERO = np.zeros(256, dtype=np.int64)

def build_frozen(train, k):
    model = {}; ctx = 0; mask = (1 << (8 * k)) - 1 if k else 0
    for sym in train:
        if ctx not in model: model[ctx] = np.zeros(256, dtype=np.int64)
        model[ctx][sym] += 1
        ctx = ((ctx << 8) | sym) & mask if k else 0
    return model

def model_ship_bytes(model):
    """Honest shipped size: frozen table serialized, then xz'd (real bytes)."""
    packed = {c: row.astype(np.uint16).tobytes() for c, row in model.items()}
    return len(lzma.compress(pickle.dumps(packed), preset=6))

xz_test = len(lzma.compress(test, preset=9))
gz_test = len(gzip.compress(test, 9))
print(f"baselines on test: raw {len(test)}  gzip {gz_test}  xz {xz_test} "
      f"({xz_test*8/len(test):.3f} b/char)\n")
print(f"{'order k':>7} {'test b/char':>11} {'tape B':>9} {'model B(xz)':>12} "
      f"{'TOTAL self':>11} {'amortized':>10} {'lossless':>9}")
print("-" * 76)

for k in [0, 1, 2, 3]:
    model = build_frozen(train, k)
    enc = encode(test, model, k)
    dec = decode(enc, model, k, len(test))
    ok = dec == test
    tape = len(enc)
    msz = model_ship_bytes(model)
    bpc = tape * 8 / len(test)
    total = msz + tape
    print(f"{k:>7} {bpc:>11.3f} {tape:>9} {msz:>12} {total:>11} {tape:>10} "
          f"{'PASS' if ok else 'FAIL':>9}")

print("\nTAPE shrinks with k (better prediction); MODEL grows with k (more contexts).")
print("TOTAL self-contained = the Hutter number (model must ship). AMORTIZED = tape")
print("only, valid ONLY when both ends already share the frozen model.")