"""
"How LLMs drop data recoverably" = superposition = compressed sensing.
Store an N-feature vector in d < N dimensions via a random projection Phi (d x N);
recover it. This is exactly what a residual stream does: pack many features into
few dims with near-orthogonal directions, read them back by correlation.

Claim under test: recovery is near-LOSSLESS iff the vector is SPARSE (few features
active). Past the compressed-sensing bound (k ~ d/(2 log N)) recovery COLLAPSES —
"recoverable" becomes "lost". So an LLM drops N->d recoverably ONLY to the extent
the signal is sparse/redundant; on dense/random data it cannot recover. Same
entropy wall, dressed as dimensionality reduction.

Two decoders:
  OMP  - greedy orthogonal matching pursuit (pick most-correlated feature, repeat);
         given the TRUE k as its budget = best case for the model.
  MF   - matched filter: x_hat = top-k of Phi^T y (the raw dot-product readout).
"""
import numpy as np
rng = np.random.default_rng(0)

N, d = 1024, 128           # 1024 features packed into a 128-dim stream (8x squeeze)
TRIALS = 20

def make_phi():
    Phi = rng.standard_normal((d, N))
    return Phi / np.linalg.norm(Phi, axis=0, keepdims=True)   # unit-norm columns

def omp(Phi, y, k):
    r = y.copy(); supp = []
    for _ in range(k):
        j = int(np.argmax(np.abs(Phi.T @ r)))
        if j in supp: break
        supp.append(j)
        Ps = Phi[:, supp]
        xs, *_ = np.linalg.lstsq(Ps, y, rcond=None)
        r = y - Ps @ xs
    x = np.zeros(N);
    if supp: x[supp] = xs
    return x, supp

def rel_err(a, b):
    return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12))

bound = d / (2 * np.log(N))
print(f"N={N} features -> d={d} dims  ({N/d:.0f}x squeeze).  "
      f"CS bound k* ~ d/(2 ln N) = {bound:.1f}\n")
print(f"{'k active':>9}{'stored N->d':>12}{'OMP rel-err':>13}{'OMP supp-exact':>15}"
      f"{'MF rel-err':>12}   recover?")
print("-" * 82)

for k in [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024]:
    oe = me = se = 0.0
    for _ in range(TRIALS):
        Phi = make_phi()
        x = np.zeros(N)
        supp_true = rng.choice(N, size=k, replace=False)
        x[supp_true] = rng.standard_normal(k)
        y = Phi @ x                              # the compressed d-dim code
        xo, supp = omp(Phi, y, k)
        oe += rel_err(xo, x)
        se += len(set(supp) & set(supp_true.tolist())) / k
        # matched filter: top-k by |correlation|, LS on that support
        corr = np.abs(Phi.T @ y)
        top = np.argsort(corr)[-k:]
        xm = np.zeros(N)
        xm[top], *_ = np.linalg.lstsq(Phi[:, top], y, rcond=None)
        me += rel_err(xm, x)
    oe/=TRIALS; me/=TRIALS; se/=TRIALS
    verdict = "LOSSLESS-ish" if oe < 1e-6 else ("degrading" if oe < 0.3 else "LOST")
    print(f"{k:>9}{d/N:>11.2f}x{oe:>13.2e}{se*100:>14.0f}%{me:>12.2e}   {verdict}")

print(f"\nSmall k: 8x squeeze recovered to ~machine-zero (sparse -> lossless-ish).")
print(f"Past k* ~ {bound:.0f}: OMP support-recovery breaks, error -> O(1) = LOST.")
print(f"k=N (dense/random): error ~1.0 — a random vector CANNOT be dropped-and-recovered.")
print(f"=> 'recoverable' is bought entirely by sparsity; it is not free compression.")