import json, numpy as np, itertools, hashlib
R=json.load(open('results_raw.json')); U=R['units']; T=R['alloc_table']
DOCS=['wikiA','wikiB','mathA']; L=8
out={'config':R['config'],'calib_err':R['calib_err']}
def m(xs): return float(np.mean(xs))
# 1. variant table per doc
tab={}
for doc in DOCS:
    for v,b in [('A',3),('B',3),('C',3),('D',3),('A',4),('B',4)]:
        u=[x for x in U if x['doc']==doc and x['var']==v and x['bits']==b]
        u7=[x for x in u if x['layer']==7]
        tab[f'{doc}|{v}{b}']=dict(bpv=u[0]['bpv'],K_err=m([x['K_err'] for x in u]),V_err=m([x['V_err'] for x in u]),
            V7_err=m([x['V_err'] for x in u7]),attn_err=m([x['attn_err'] for x in u]),KL=m([x['KL'] for x in u]),
            attn_err_Konly=m([x['attn_err_Konly'] for x in u]),KL_Konly=m([x['KL_Konly'] for x in u]),
            attn_err_L7=m([x['attn_err'] for x in u7]))
out['variant_table']=tab
# 2. D vs B per unit, and ranking per doc
cmp={}
for doc in DOCS:
    g=lambda v,k:{(x['layer'],x['head']):x[k] for x in U if x['doc']==doc and x['var']==v and x['bits']==3}
    for k in ('K_err','V_err','attn_err','KL','attn_err_Konly'):
        b,d=g('B',k),g('D',k); cmp[f'{doc}|{k}']=dict(D_beats_B=sum(d[i]<b[i] for i in b),units=len(b),mean_ratio_D_over_B=m([d[i]/b[i] for i in b]))
    cmp[f'{doc}|rank_attn_err']=sorted('ABCD',key=lambda v:tab[f'{doc}|{v}3']['attn_err'])
    cmp[f'{doc}|rank_K_err']=sorted('ABCD',key=lambda v:tab[f'{doc}|{v}3']['K_err'])
out['D_vs_B']=cmp
# 3. allocation: DP over layers, fit on wikiA, evaluate on all docs
OPT=['2,32','3,32','4,32','2,128','3,128','4,128']
bpv=lambda o:int(o.split(',')[0])+16/int(o.split(',')[1])
alloc={}
for v in ('B','D'):
    for budget in (3.2,):
        # DP on sum of bpv*8 (integer units of 1/8 bit)
        best={0:(0.0,[])}
        for l in range(L):
            nb={}
            for s,(e,ch) in best.items():
                for ok,ov in itertools.product(OPT,OPT):
                    c=s+round(8*(bpv(ok)+bpv(ov)))
                    if c>round(8*budget*2*L): continue
                    ee=e+T[f'wikiA|{v}|{l}|{ok}|{ov}'][0]
                    if c not in nb or nb[c][0]>ee: nb[c]=(ee,ch+[(ok,ov)])
            best=nb
        e,ch=min(best.values(),key=lambda t:t[0])
        avg=m([bpv(o) for p in ch for o in p])
        ev={doc:dict(attn_err=m([T[f'{doc}|{v}|{l}|{ch[l][0]}|{ch[l][1]}'][0] for l in range(L)]),
                     KL=m([T[f'{doc}|{v}|{l}|{ch[l][0]}|{ch[l][1]}'][1] for l in range(L)]),
                     ref_rot_q4_0_attn_err=m([T[f'{doc}|B|{l}|4,32|4,32'][0] for l in range(L)]),
                     ref_rot_q4_0_KL=m([T[f'{doc}|B|{l}|4,32|4,32'][1] for l in range(L)]),
                     uniform_3b_g32_attn_err=m([T[f'{doc}|{v}|{l}|3,32|3,32'][0] for l in range(L)]),
                     uniform_3b_g128_attn_err=m([T[f'{doc}|{v}|{l}|3,128|3,128'][0] for l in range(L)]),
                     uniform_3b_g128_KL=m([T[f'{doc}|{v}|{l}|3,128|3,128'][1] for l in range(L)])) for doc in DOCS}
        alloc[f'{v}|{budget}']=dict(fit_on='wikiA',avg_bpv=avg,per_layer=[dict(layer=l,K=ch[l][0],V=ch[l][1]) for l in range(L)],eval=ev)
out['allocation']=alloc
s=json.dumps(out,indent=1,sort_keys=True); open('results.json','w').write(s)
print(s)