# E21u: bit-level analysis of real KV cache captured in E21t (Qwen3.5-4B, 8 full-attn layers, 4 KV heads x 256 dims)
import numpy as np, struct, mmap, json, sys, time
D='C:/AGI/agi-research/e21_budget/e21t/'
OUT='C:/AGI/agi-research/e21_budget/e21u/'
import os; os.makedirs(OUT,exist_ok=True)
docs={'wikiA':28952,'wikiB':27083,'mathA':30287}
H,HD=4,256
def blocks(fn,N,t_expect,R):
    f=open(fn,'rb'); mm=mmap.mmap(f.fileno(),0,access=mmap.ACCESS_READ)
    o=mm.find(struct.pack('<iQ',t_expect,R),0,8_000_000); out=[]
    for i in range(16):
        t,r=struct.unpack_from('<iQ',mm,o); assert t==t_expect and r==R,(i,t,r)
        out.append(o+12); o=o+12+N*r
    return mm,out
def layer_f16(mm,off,N): return np.frombuffer(mm,dtype=np.float16,count=N*1024,offset=off).reshape(N,1024).astype(np.float32)
def layer_q40(mm,off,N):
    b=np.frombuffer(mm,dtype=np.uint8,count=N*576,offset=off).reshape(N*32,18)
    d=b[:,:2].copy().view(np.float16).astype(np.float32)
    q=b[:,2:]; lo=(q&15).astype(np.int8)-8; hi=(q>>4).astype(np.int8)-8
    v=np.concatenate([lo,hi],1).astype(np.float32)*d
    return v.reshape(N,1024)
def ent16(x16):
    h=np.bincount(x16.view(np.uint16).ravel(),minlength=65536).astype(np.float64); p=h[h>0]/h.sum(); return float(-(p*np.log2(p)).sum())
def had(n):
    Hm=np.array([[1.]]);
    while Hm.shape[0]<n: Hm=np.block([[Hm,Hm],[Hm,-Hm]])
    return Hm/np.sqrt(n)
HADM=had(HD).astype(np.float32)
def q_tokgroup(x,bits,g=32):
    s=x.reshape(-1,g); m=np.abs(s).max(1,keepdims=True); L=2**(bits-1)-1 if bits>1 else 1
    L=max(L,1); sc=m/ (2**(bits-1)) + 1e-12
    q=np.clip(np.round(s/sc),-(2**(bits-1)),2**(bits-1)-1)*sc; return q.reshape(x.shape)
def q_minmax_tok(x,bits,g=32):
    s=x.reshape(-1,g); mn=s.min(1,keepdims=True); mx=s.max(1,keepdims=True); sc=(mx-mn)/(2**bits-1)+1e-12
    return (np.round((s-mn)/sc)*sc+mn).reshape(x.shape)
def q_minmax_chan(x,bits,g=32):
    N=x.shape[0]//g*g; y=x.copy(); s=x[:N].reshape(-1,g,x.shape[1])
    mn=s.min(1,keepdims=True); mx=s.max(1,keepdims=True); sc=(mx-mn)/(2**bits-1)+1e-12
    y[:N]=(np.round((s-mn)/sc)*sc+mn).reshape(N,x.shape[1]); return y
def rel(a,b): return float(((a-b)**2).sum()/((b-b.mean(0))**2).sum())
rng=np.random.default_rng(0)
res={}; chanstats={}; bases={}
t0=time.time()
for doc,N in docs.items():
    mm,offs=blocks(D+f'kv_f16_{doc}.bin',N,1,2048)
    for li,off in enumerate(offs):
        kind='K' if li<8 else 'V'; L=li%8; tag=f'{doc}_{kind}{L}'
        x16=np.frombuffer(mm,dtype=np.float16,count=N*1024,offset=off).reshape(N,1024)
        x=x16.astype(np.float32)
        r={'doc':doc,'kind':kind,'layer':L,'n':N}
        r['std']=float(x.std()); r['absmax']=float(np.abs(x).max())
        cam=np.abs(x).max(0); r['outlier_chan_frac']=float((cam>5*np.median(cam)).mean())
        z=(x-x.mean(0))/ (x.std(0)+1e-9); r['kurt_mean']=float((z**4).mean())
        r['entropy_f16_bits']=ent16(x16)
        # quant error (relative to per-channel variance), 4096-token slice for speed
        idx=slice(0,min(N,8192)); xs=x[idx]
        xh=(xs.reshape(-1,H,HD)@HADM).reshape(xs.shape)
        for b in (2,3,4):
            r[f'relerr_tok_absmax_{b}b']=rel(q_tokgroup(xs,b),xs)
            r[f'relerr_tok_minmax_{b}b']=rel(q_minmax_tok(xs,b),xs)
            r[f'relerr_chan_minmax_{b}b']=rel(q_minmax_chan(xs,b),xs)
            xq=(q_minmax_tok(xh,b).reshape(-1,H,HD)@HADM.T).reshape(xs.shape)
            r[f'relerr_hadamard_tok_minmax_{b}b']=rel(xq,xs)
        # low rank per head on sampled tokens
        sm=x[rng.choice(N,min(N,6000),replace=False)]
        k90=[];k99=[];B=[]
        for h in range(H):
            a=sm[:,h*HD:(h+1)*HD]; a=a-a.mean(0); u,s,vt=np.linalg.svd(a,full_matrices=False); e=np.cumsum(s**2)/np.sum(s**2)
            k90.append(int(np.searchsorted(e,0.90)+1)); k99.append(int(np.searchsorted(e,0.99)+1)); B.append(vt[:32])
        r['pca_k90_per_head']=k90; r['pca_k99_per_head']=k99
        bases[tag]=B
        # temporal
        a=x[1:]; b_=x[:-1]
        r['adj_cos']=float(((a*b_).sum(1)/(np.linalg.norm(a,axis=1)*np.linalg.norm(b_,axis=1)+1e-9)).mean())
        r['delta_var_ratio']=float(((a-b_)**2).sum()/((x-x.mean(0))**2).sum())
        chanstats[tag]=(x.mean(0),x.std(0))
        res[tag]=r
        print(tag,round(time.time()-t0),json.dumps({k:v for k,v in r.items() if k in('entropy_f16_bits','relerr_tok_absmax_3b','relerr_hadamard_tok_minmax_3b','pca_k90_per_head','adj_cos')}),flush=True)
# replication across docs
rep={}
for kind in 'KV':
    for L in range(8):
        for d1,d2 in (('wikiA','wikiB'),('wikiA','mathA'),('wikiB','mathA')):
            m1,s1=chanstats[f'{d1}_{kind}{L}']; m2,s2=chanstats[f'{d2}_{kind}{L}']
            ov=[]
            for h in range(H):
                B1=bases[f'{d1}_{kind}{L}'][h]; B2=bases[f'{d2}_{kind}{L}'][h]; ov.append(float((np.linalg.svd(B1@B2.T,compute_uv=False)**2).mean()))
            rep[f'{kind}{L}_{d1}_{d2}']={'chan_mean_corr':float(np.corrcoef(m1,m2)[0,1]),'chan_std_corr':float(np.corrcoef(s1,s2)[0,1]),'top32_subspace_overlap':ov}
# q4_0 stored vs f16 (wikiA)
mmq,offq=blocks(D+'kv_q4_0_wikiA.bin',28952,2,576); mmf,offf=blocks(D+'kv_f16_wikiA.bin',28952,1,2048)
q4={}
for li in range(16):
    xf=layer_f16(mmf,offf[li],28952); xq=layer_q40(mmq,offq[li],28952)
    q4[('K' if li<8 else 'V')+str(li%8)]=rel(xq,xf)
json.dump({'per_layer':res,'replication':rep,'q4_0_stored_relerr_vs_f16':q4},open(OUT+'results.json','w'),indent=1)
print('Q4STORED',json.dumps(q4)); print('=== E21uDONE ===',flush=True)