"""CPU-only observed activation sparsity; no checkpoint changes or training."""
import json,csv,time,resource,hashlib
from pathlib import Path
import numpy as np
from tokenizers import Tokenizer
from functional.model import StreamingLM,silu,sigmoid


def metrics(x):
    x=np.asarray(x,dtype=np.float64).reshape(-1,x.shape[-1]); a=np.abs(x); energy=np.square(x).sum(0); total=energy.sum()
    rms=float(np.sqrt(np.mean(x*x))); cutoff=.01*rms
    ordered=np.sort(energy)[::-1]; n90=int(np.searchsorted(np.cumsum(ordered),.9*total)+1) if total else 0
    return dict(elements=x.size,channels=x.shape[-1],rms=rms,zero_rate=float(np.mean(x==0)),negative_rate=float(np.mean(x<0)),near_001_rms=float(np.mean(a<=cutoff)),near_01_rms=float(np.mean(a<=.1*rms)),ever_active_001=float(np.mean(np.any(a>cutoff,axis=0))),channels_for_90pct_energy=n90,energy90_channel_fraction=n90/x.shape[-1])

class Probe(StreamingLM):
    def record(self,x,name):
        for width in (0,64,128):
            for start in range(0,x.shape[-1],width or x.shape[-1]):
                stop=min(start+(width or x.shape[-1]),x.shape[-1]); z=x[...,start:stop]
                self.rows.append(dict(sample=self.sample,module=name,width=width,start=start,stop=stop,**metrics(z)))
    def linear(self,x,name,rows=256):
        y=super().linear(x,name,rows); self.record(y,name+'.output')
        if name.endswith('self_attn.q_proj'):
            dim=self.config['head_dim']; mixed=y.reshape(*y.shape[:-1],-1,2*dim)
            self.record(sigmoid(mixed[...,dim:]).reshape(*y.shape[:-1],-1),name+'.sigmoid_gate')
        return y
    def mlp(self,h,layer):
        p=f'layers.{layer}.mlp.'; pre=self.linear(h,p+'gate_proj'); g=silu(pre); u=self.linear(h,p+'up_proj')
        self.record(g,p+'silu_gate'); self.record(g*u,p+'gated_product')
        return self.linear(g*u,p+'down_proj')

def main():
    import argparse
    p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--out',type=Path,required=True);a=p.parse_args();a.out.mkdir(parents=True,exist_ok=False)
    started=time.monotonic();cfg=json.loads((a.source/'model_config.json').read_text()); lm=Probe(a.source/'weights',cfg,seconds=3600);lm.rows=[]
    samples=json.loads(Path('crossqkv/fresh_samples.json').read_text())['samples'];tok=Tokenizer.from_file(str(a.source/'tokenizer.json'))
    for s in samples:
        lm.sample=s['id']; ids=tok.encode(s['text']).ids[:32];s['token_ids']=ids
        lm.forward(np.array([ids],dtype=np.int64))
        first=not (a.out/'metrics.csv').exists()
        with (a.out/'metrics.csv').open('a') as f:
            w=csv.DictWriter(f,fieldnames=list(lm.rows[0]))
            if first:w.writeheader()
            w.writerows(lm.rows)
        print('completed',s['id'],len(lm.rows),flush=True);lm.rows.clear()
    summary={'completed':True,'elapsed_seconds':time.monotonic()-started,'peak_rss_kib':resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,'config_activation':lm.config.get('hidden_act'),'layer_count':len(lm.config['layer_types']),'samples':samples,'rows':sum(1 for _ in (a.out/'metrics.csv').open())-1,'block_widths':[64,128],'threshold_reference':'RMS within each sample/module/block, thresholds 0.01 and 0.1; exact zero measured separately','training':False,'weights_deleted':False,'limits':['4 short texts, 32 tokens each, descriptive only','Block is contiguous output channels, not Transformer block','Small output does not prove removable computation or parameter interchangeability','All samples previously used; no generalization claim','Relative threshold differs between blocks; width=0 is whole module reference']}
    (a.out/'summary.json').write_text(json.dumps(summary,indent=2))
if __name__=='__main__':main()