"""
Decisive test: can a low-rank / spectral summary of the order-1 byte
transition matrix beat the raw empirical counts ("KV cache") at prediction?

Claim: NO. The empirical conditional distribution is the MLE — it MINIMIZES
cross-entropy on its own data. Any low-rank reconstruction (the regime the
characteristic polynomial / eigen-summary lives in) has cross-entropy >= the
empirical order-1 entropy. So the polynomial cannot add predictive information
over the count matrix it is derived from. It can only match (full rank) or hurt.

Also report order-0 and order-2 to show the REAL lever is context DEPTH, not rank.
"""
import os, sys, math
import numpy as np

def load_corpus(root, cap=1_500_000, max_files=3000):
    buf = bytearray()
    seen = 0
    for dirpath, _, files in os.walk(root):
        for fn in files:
            if not (fn.endswith(".md") or fn.endswith(".txt")):
                continue
            seen += 1
            if seen > max_files:
                break
            try:
                with open(os.path.join(dirpath, fn), "rb") as f:
                    buf += f.read()
            except Exception:
                pass
            if len(buf) >= cap:
                return bytes(buf[:cap])
        if seen > max_files or len(buf) >= cap:
            break
    return bytes(buf[:cap])

def order0_bits(data):
    c = np.bincount(np.frombuffer(data, dtype=np.uint8), minlength=256).astype(float)
    p = c / c.sum()
    nz = p > 0
    return float(-(p[nz] * np.log2(p[nz])).sum())

def order1_counts(data):
    a = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
    idx = a[:-1] * 256 + a[1:]
    C = np.bincount(idx, minlength=256*256).reshape(256, 256).astype(float)
    return C

def cond_xent(C, Q):
    # cross-entropy of the data (weighted by true counts C) under model Q(j|i)
    N = C.sum()
    Q = np.clip(Q, 1e-12, None)
    return float(-(C * np.log2(Q)).sum() / N)

def order1_empirical_bits(C):
    rs = C.sum(axis=1, keepdims=True)
    P = C / np.clip(rs, 1e-12, None)
    return cond_xent(C, P), P

def lowrank_bits(C, k):
    # rank-k approximation of the conditional matrix, clipped & renormalized
    rs = C.sum(axis=1, keepdims=True)
    P = C / np.clip(rs, 1e-12, None)
    U, S, Vt = np.linalg.svd(P, full_matrices=False)
    Pk = (U[:, :k] * S[:k]) @ Vt[:k]
    Pk = np.clip(Pk, 0, None)
    Pk = Pk / np.clip(Pk.sum(axis=1, keepdims=True), 1e-12, None)
    return cond_xent(C, Pk)

def order2_bits(data):
    a = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
    if len(a) < 3:
        return float("nan")
    ctx = a[:-2] * 256 + a[1:-2+1]  # pairs (b0,b1)
    ctx = a[:-2] * 256 + a[1:-1]
    nxt = a[2:]
    from collections import defaultdict
    cnt = defaultdict(lambda: np.zeros(256))
    for c, n in zip(ctx.tolist(), nxt.tolist()):
        cnt[c][n] += 1
    N = len(nxt)
    H = 0.0
    for c, row in cnt.items():
        rs = row.sum()
        p = row / rs
        nz = p > 0
        H += (rs / N) * float(-(p[nz] * np.log2(p[nz])).sum())
    return H

def main():
    root = sys.argv[1] if len(sys.argv) > 1 else "/home/allaun/Research Stack/6-Documentation"
    data = load_corpus(root)
    print(f"corpus: {len(data)} bytes from {root}\n")
    h0 = order0_bits(data)
    C = order1_counts(data)
    h1, _ = order1_empirical_bits(C)
    print(f"order-0 entropy                : {h0:6.3f} bits/byte")
    print(f"order-1 empirical (KV cache)   : {h1:6.3f} bits/byte   <-- MLE optimum for order-1")
    print(f"order-2 empirical (deeper ctx) : {order2_bits(data):6.3f} bits/byte   <-- context DEPTH helps\n")
    print("rank-k reconstruction of the order-1 conditional matrix")
    print("(the regime the char-poly / eigen-summary lives in):")
    for k in [1, 2, 4, 8, 16, 32, 64, 128, 255, 256]:
        hk = lowrank_bits(C, k)
        flag = "  >= order-1 (worse or equal)" if hk >= h1 - 1e-9 else "  << BEATS order-1 (!)"
        print(f"  rank {k:3d} : {hk:6.3f} bits/byte{flag}")

main()