import json, numpy as np
from spiderweb_experiment import SpiderwebAttention, dense_mean, local_mean, hierarchy_work
n,d=256,8
x=np.zeros((n,d)); x[0,0]=1.
sp=SpiderwebAttention(d,levels=8,radius=1)
y_sp,_=sp(x); y_de=dense_mean(x); y_lo=local_mean(x,2)
res={}
for name,y in [('dense',y_de),('local',y_lo),('spiderweb',y_sp)]:
    res[name]={'token_1':float(y[1,0]),'token_4':float(y[4,0]),
               'token_128':float(y[128,0]),'token_255':float(y[255,0]),
               'nonzero_tokens':int(np.count_nonzero(np.abs(y[:,0])>1e-12))}
res['work']=hierarchy_work(n)[0]
with open('impulse_results.json','w') as f: json.dump(res,f,indent=2)
print(json.dumps(res,indent=2))