"""Find repeated learned scalar MLP functions on real hidden inputs, no training.

Four fixed layers; calibration-only selection; one frozen pair; independent heldout
functional test and two actual network substitutions. No scale fitting or R.
"""
import argparse
import hashlib
import json
from pathlib import Path
import resource
import time
import numpy as np
from .model import StreamingLM,silu
from .run import write_json,save_csv

LAYERS=[5,6,20,21]


def pair_metrics(x,y):
    x,y=x.astype(np.float64),y.astype(np.float64)
    xx,yy=float(x@x),float(y@y);delta=float((x-y)@(x-y))
    xc,yc=x-x.mean(),y-y.mean()
    return {'relative_l2_a_to_b':float(np.sqrt(delta/max(xx,1e-30))),
            'relative_l2_b_to_a':float(np.sqrt(delta/max(yy,1e-30))),
            'cosine':float(x@y/max(np.sqrt(xx*yy),1e-30)),
            'pearson':float(xc@yc/max(np.linalg.norm(xc)*np.linalg.norm(yc),1e-30)),
            'rms_a':float(np.sqrt(xx/x.size)),'rms_b':float(np.sqrt(yy/y.size)),
            'std_over_rms_a':float(np.std(x)/max(np.sqrt(xx/x.size),1e-30)),
            'std_over_rms_b':float(np.std(y)/max(np.sqrt(yy/y.size),1e-30)),
            'effective_energy_tokens_a':float(xx*xx/max(float((x**4).sum()),1e-30)),
            'effective_energy_tokens_b':float(yy*yy/max(float((y**4).sum()),1e-30))}


class NeuronReuseLM(StreamingLM):
    def configure(self,layer,target,donor):self.reused_layer,self.target,self.donor=layer,target,donor
    def mlp(self,h,layer):
        p=f'layers.{layer}.mlp.'
        z=silu(self.linear(h,p+'gate_proj'))*self.linear(h,p+'up_proj')
        if layer==self.reused_layer:z[...,self.target]=z[...,self.donor]
        return self.linear(z,p+'down_proj')


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);a=p.parse_args();root=a.root
    protocol={'layers':LAYERS,'unit':'post_silu_gate_times_up_scalar_function; keep each output column',
              'thresholds':[.05,.10,.20],'maximum_directional_relative_l2':.20,'cosine_min':.98,'pearson_min':.95,
              'rms_floor':'at least within-layer median calibration neuron RMS','std_over_rms_min':.25,
              'effective_energy_tokens_min':16,'max_selected_pairs':1,'selection':'smallest worst directional calibration error; then layer/id',
              'heldout_per_text_error_max':.35,'heldout_quality_mean_nll_delta_max':.01,'heldout_quality_sequence_p95_delta_max':.05,
              'training':False,'scale_fitting':False,'cost_required':False}
    protocol_path=root/'neuron-reuse-protocol.json'
    if protocol_path.exists():
        old=json.loads(protocol_path.read_text());assert all(old[k]==v for k,v in protocol.items())
    else:write_json(protocol_path,{**protocol,'registered_unix':time.time()})
    ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text());remaining=int(7200-ledger['spent_seconds']-30)
    if remaining<120:raise TimeoutError('Remaining cumulative budget insufficient')
    started=time.monotonic();result={'state':'running','selected':None,'function_pass':False,'interventions':[]}
    try:
        cfg=json.loads((root/'model_config.json').read_text());lm=NeuronReuseLM(root/'weights',cfg,seconds=remaining)
        data=np.load(root/'baseline-activations.npz');samples=json.loads((root/'samples-tokenized.json').read_text())['samples']
        cal=np.array([i for i,s in enumerate(samples) if s['split']=='calibration']);held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout'])
        activations={};summaries=[];best=None
        for layer in LAYERS:
            h=data[f'mlp_{layer}'];pfx=f'layers.{layer}.mlp.'
            z=silu(lm.linear(h,pfx+'gate_proj'))*lm.linear(h,pfx+'up_proj');activations[layer]=z
            # Profiles use all calibration tokens; never use heldout in selection.
            x=np.ascontiguousarray(z[cal].reshape(-1,z.shape[-1]).T,dtype=np.float64)
            n=x.shape[1];normsq=(x*x).sum(1);means=x.mean(1);variance=np.maximum(normsq-n*means**2,0)
            rms=np.sqrt(normsq/n);floor=float(np.median(rms));effective=normsq**2/np.maximum((x**4).sum(1),1e-30)
            active=(rms>=floor)&(np.sqrt(variance/np.maximum(normsq,1e-30))>=.25)&(effective>=16)
            counts={str(t):0 for t in protocol['thresholds']};tested=0;layerbest=None
            for start in range(0,x.shape[0],256):
                lm.guard();end=min(start+256,x.shape[0]);dot=x[start:end]@x.T
                delta=np.maximum(normsq[start:end,None]+normsq[None,:]-2*dot,0)
                error=np.sqrt(delta/np.maximum(np.minimum(normsq[start:end,None],normsq[None,:]),1e-30))
                cosine=dot/np.maximum(np.sqrt(normsq[start:end,None]*normsq[None,:]),1e-30)
                pearson=(dot-n*means[start:end,None]*means[None,:])/np.maximum(np.sqrt(variance[start:end,None]*variance[None,:]),1e-30)
                unique=np.arange(start,end)[:,None]<np.arange(x.shape[0])[None,:]
                allowed=unique&active[start:end,None]&active[None,:]&(cosine>=.98)&(pearson>=.95)
                tested+=int((unique&active[start:end,None]&active[None,:]).sum())
                for threshold in protocol['thresholds']:counts[str(threshold)]+=int((allowed&(error<=threshold)).sum())
                mask=allowed&(error<=.20)
                if mask.any():
                    score=np.where(mask,error,np.inf);row,col=np.unravel_index(np.argmin(score),score.shape)
                    candidate={'layer':layer,'a':int(start+row),'b':int(col),'calibration_worst_relative_l2':float(score[row,col])}
                    key=lambda c:(c['calibration_worst_relative_l2'],c['layer'],c['a'],c['b'])
                    if layerbest is None or key(candidate)<key(layerbest):layerbest=candidate
                    if best is None or key(candidate)<key(best):best=candidate
            summaries.append({'layer':layer,'neurons':int(x.shape[0]),'eligible_neurons':int(active.sum()),'eligible_unique_pairs_checked':tested,'rms_floor':floor,'calibration_pair_counts':counts,'best':layerbest})
            write_json(root/'neuron-reuse-search-summary.json',summaries)
            print(json.dumps(summaries[-1]),flush=True)
        np.savez(root/'neuron-reuse-activations.npz',**{f'layer{k}':v for k,v in activations.items()})
        write_json(root/'neuron-reuse-selection.json',{'selected':best,'selection_uses_heldout':False})
        result['selected']=best
        if best is not None:
            layer,i,j=best['layer'],best['a'],best['b'];z=activations[layer]
            metrics=[]
            for split,index in [('calibration',cal),('heldout',held)]:
                metrics.append({'split':split,'scope':'pooled',**pair_metrics(z[index,:,i].reshape(-1),z[index,:,j].reshape(-1))})
                for idx in index:metrics.append({'split':split,'scope':samples[idx]['id'],**pair_metrics(z[idx,:,i],z[idx,:,j])})
            save_csv(root/'neuron-reuse-function-metrics.csv',metrics)
            pooled=next(r for r in metrics if r['split']=='heldout' and r['scope']=='pooled')
            floor=next(r['rms_floor'] for r in summaries if r['layer']==layer)
            function_pass=(max(pooled['relative_l2_a_to_b'],pooled['relative_l2_b_to_a'])<=.20 and pooled['cosine']>=.98 and pooled['pearson']>=.95 and
                           min(pooled['std_over_rms_a'],pooled['std_over_rms_b'])>=.25 and min(pooled['effective_energy_tokens_a'],pooled['effective_energy_tokens_b'])>=16 and
                           min(pooled['rms_a'],pooled['rms_b'])>=floor and
                           all(max(r['relative_l2_a_to_b'],r['relative_l2_b_to_a'])<=.35 for r in metrics if r['split']=='heldout' and r['scope']!='pooled'))
            result.update(function_pass=bool(function_pass),heldout_function_metrics=pooled)
            # Show parameter similarity separately from actual functional similarity.
            wd=[]
            for role in ('gate_proj','up_proj'):
                t=lm.tensors[f'model.language_model.layers.{layer}.mlp.{role}.weight'];d=t.shape[1]
                wa,wb=t.read(slice(i*d,(i+1)*d)),t.read(slice(j*d,(j+1)*d))
                wd.append({'role':role,'weight_cosine':float(wa.astype(np.float64)@wb/max(np.linalg.norm(wa.astype(np.float64))*np.linalg.norm(wb.astype(np.float64)),1e-30))});t.evict()
            result['weight_diagnostics']=wd
            if function_pass:
                ids=np.array([s['token_ids'] for s in samples],dtype=np.int64);base=np.load(root/'baseline-heldout-token-nll.npy').mean(1)
                for target,donor in ((i,j),(j,i)):
                    lm.configure(layer,target,donor);h,_=lm.forward(ids[held]);nll=lm.nll(h,ids[held])
                    np.save(root/f'neuron-reuse-{target}-uses-{donor}-token-nll.npy',nll)
                    delta=nll.mean(1)-base
                    result['interventions'].append({'target':target,'donor':donor,'sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_delta':float(np.quantile(delta,.95)),'passed':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)})
                    write_json(root/'neuron-reuse-summary.json',result)
            # No fallback after looking at heldout.
        result.update(state='completed',empirical_repeated_function_positive=bool(result['function_pass'] and len(result['interventions'])==2 and all(r['passed'] for r in result['interventions'])),
                      claim_limit='Distinct learned scalar MLP subfunctions on the specified short-text distribution; not whole-layer duplication, global equivalence or broad task-capability proof. No performance requirement.')
    except Exception as exc:result.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise
    finally:
        elapsed=time.monotonic()-started;result.update(elapsed_seconds=elapsed,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
        ledger=json.loads(ledger_path.read_text());ledger['spent_seconds']+=elapsed;ledger['calls'].append({'module':'functional.neuron_reuse','seconds':elapsed});write_json(ledger_path,ledger)
        write_json(root/'neuron-reuse-summary.json',result)


if __name__=='__main__':main()