"""Render scientific figures from recorded CPU measurements."""
import csv,json
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import sys
root=Path(sys.argv[1]); rows=list(csv.DictReader((root/'metrics.csv').open()))
full=[r for r in rows if r['width']=='0']
fig,axs=plt.subplots(1,3,figsize=(15,4.5),layout='constrained')
for suffix,label in [('gate_proj.output','gate preactivation'),('silu_gate','SiLU gate'),('gated_product','gated product')]:
    selected=[r for r in full if r['module'].endswith('mlp.'+suffix)]
    for ax,key in zip(axs,['zero_rate','near_001_rms','near_01_rms']):
        xs=[];ys=[]
        for layer in range(24):
            z=[float(r[key]) for r in selected if r['module'].startswith(f'layers.{layer}.')];xs.append(layer);ys.append(np.mean(z))
        ax.plot(xs,ys,label=label)
for ax,title in zip(axs,['Exact zero fraction','Near zero: |x| <= 0.01 RMS','Near zero: |x| <= 0.1 RMS']):
    ax.set(title=title,xlabel='Layer index',ylabel='Fraction');ax.set_ylim(bottom=0);ax.grid(alpha=.2)
axs[-1].legend();fig.suptitle('Observed MLP activation sparsity: four texts, 32 tokens each; no training')
root.joinpath('figures').mkdir(exist_ok=True)
for ext in ('png','pdf','svg'):fig.savefig(root/'figures'/('mlp_activation.'+ext),dpi=180)
plt.close(fig)
blocks=[r for r in rows if r['width']=='128' and r['module'].endswith('mlp.gated_product')]
fig,axs=plt.subplots(1,2,figsize=(12,4.5),layout='constrained')
for ax,key,title in zip(axs,['near_01_rms','energy90_channel_fraction'],['Near-zero fraction at 0.1 block RMS','Fraction of channels carrying 90% energy']):
    matrix=[]
    for layer in range(24):
        rr=[r for r in blocks if r['module'].startswith(f'layers.{layer}.')]; starts=sorted(set(int(r['start']) for r in rr));matrix.append([np.mean([float(r[key]) for r in rr if int(r['start'])==s]) for s in starts])
    im=ax.imshow(matrix,aspect='auto',vmin=0,vmax=1);ax.set(title=title,xlabel='Contiguous 128-channel block',ylabel='Layer index');fig.colorbar(im,ax=ax)
fig.suptitle('MLP gated-product blocks: average over four observed texts')
for ext in ('png','pdf','svg'):fig.savefig(root/'figures'/('block_activity.'+ext),dpi=180)
agg={}
for suffix in ['mlp.gate_proj.output','mlp.silu_gate','mlp.gated_product','self_attn.q_proj.output','self_attn.k_proj.output','self_attn.v_proj.output']:
    rr=[r for r in full if r['module'].endswith(suffix)]
    agg[suffix]={'observations':len(rr),**{key:{'min':min(float(r[key]) for r in rr),'mean':float(np.mean([float(r[key]) for r in rr])),'max':max(float(r[key]) for r in rr)} for key in ['zero_rate','negative_rate','near_001_rms','near_01_rms','ever_active_001','energy90_channel_fraction']}}
(root/'aggregates.json').write_text(json.dumps(agg,indent=2))