import time, json
import numpy as np
from spiderweb_experiment import SpiderwebAttention, dense_mean, local_mean, hierarchy_work

rng=np.random.default_rng(7); n=256; d=32
x=rng.normal(size=(n,d)); x[0,0]=10.0
# signal recovery: how much token n-1 receives from source coordinate 0
methods={
 'dense': lambda z: dense_mean(z),
 'local': lambda z: local_mean(z, window=2),
 'spiderweb': lambda z: SpiderwebAttention(d, levels=8, radius=1, seed=0)(z)[0],
}
out={}
for name, fn in methods.items():
    fn(x) # warmup
    t0=time.perf_counter(); reps=10
    for _ in range(reps): y=fn(x)
    sec=(time.perf_counter()-t0)/reps
    out[name]={'seconds_per_call':sec,
               'far_token_signal_dim0':float(y[-1,0]),
               'mean_abs_output':float(np.mean(np.abs(y)))}
out['counts']={'dense_pair_scores':n*n, 'local_pair_scores':n*5,
               'spiderweb_estimated_work':hierarchy_work(n)[0],
               'spiderweb_horizontal_plus_broadcast':hierarchy_work(n)[2:]}
with open('benchmark_results.json','w') as f: json.dump(out,f,indent=2)
print(json.dumps(out,indent=2))