"""
Static empirical entropy of enwik8 at order 0/1/2 (full 100MB, chunked numpy).
This is the bits/byte an *adaptive* arithmetic coder converges to (model learned
online => nothing to transmit; on 100MB the learning cost is negligible, so
static conditional entropy is a faithful estimate of the coder's output rate).
"""
import sys, numpy as np

PATH = sys.argv[1] if len(sys.argv) > 1 else "enwik8"
N = 100_000_000
CHUNK = 5_000_000

def entropy_from_counts(counts):
    counts = counts.astype(np.float64)
    tot = counts.sum()
    nz = counts > 0
    p = counts[nz] / tot
    return float(-(p * np.log2(p)).sum())

# order-0
c0 = np.zeros(256, dtype=np.int64)
# order-1 conditional: counts[prev*256+cur]
c1 = np.zeros(256*256, dtype=np.int64)
# order-2 conditional: counts[prev2*65536+prev1*256+cur]  (2^24 bins)
c2 = np.zeros(1 << 24, dtype=np.int64)

with open(PATH, "rb") as f:
    prev = np.zeros(0, dtype=np.int64)
    off = 0
    while True:
        b = f.read(CHUNK)
        if not b:
            break
        a = np.frombuffer(b, dtype=np.uint8).astype(np.int64)
        c0 += np.bincount(a, minlength=256)
        # stitch 2 bytes of previous chunk for cross-boundary bigrams/trigrams
        s = np.concatenate([prev, a])
        # order-1 pairs within s
        idx1 = s[:-1] * 256 + s[1:]
        c1 += np.bincount(idx1, minlength=256*256)
        if len(s) >= 3:
            idx2 = s[:-2] * 65536 + s[1:-1] * 256 + s[2:]
            c2 += np.bincount(idx2, minlength=1 << 24)
        prev = s[-2:]  # carry last 2 bytes

H0 = entropy_from_counts(c0)

# order-1 conditional entropy = H(cur|prev) = weighted avg of row entropies
C1 = c1.reshape(256, 256).astype(np.float64)
rs1 = C1.sum(axis=1)
tot1 = rs1.sum()
H1 = 0.0
for i in range(256):
    if rs1[i] > 0:
        p = C1[i][C1[i] > 0] / rs1[i]
        H1 += (rs1[i] / tot1) * float(-(p * np.log2(p)).sum())

# order-2 conditional entropy, streamed over the 65536 context rows
C2 = c2.reshape(65536, 256)
rs2 = C2.sum(axis=1).astype(np.float64)
tot2 = rs2.sum()
H2 = 0.0
nzc = np.where(rs2 > 0)[0]
for i in nzc:
    row = C2[i].astype(np.float64)
    row = row[row > 0]
    p = row / rs2[i]
    H2 += (rs2[i] / tot2) * float(-(p * np.log2(p)).sum())

def proj(bits):
    return bits * N / 8 / 1_048_576  # MB if we coded all 100M bytes at this rate

print(f"file: {PATH}  bytes: {N}")
print(f"order-0 entropy : {H0:6.3f} bits/byte  -> ~{proj(H0):7.2f} MB")
print(f"order-1 entropy : {H1:6.3f} bits/byte  -> ~{proj(H1):7.2f} MB")
print(f"order-2 entropy : {H2:6.3f} bits/byte  -> ~{proj(H2):7.2f} MB   (adaptive AC target)")
print("DONE_ENTROPY")