"""All adjacent layers, input-dependent influence distributions; CPU only."""
import argparse,json,time,resource,hashlib,traceback
from pathlib import Path
import numpy as np
from tokenizers import Tokenizer
from functional.model import StreamingLM
from .core import InfluenceLM,concentration,masks_from_influences,js_rows


def dump(path,data):path.write_text(json.dumps(data,indent=2,allow_nan=False)+'\n')


def distances(base,changed):
    out={}
    for key in ('Q','K','V'):
        a=base[key].astype(np.float64);b=changed[key].astype(np.float64)
        out[key]=np.linalg.norm(b-a,axis=-1)/np.maximum(np.linalg.norm(a,axis=-1),1e-12)
    out['QKV']=np.sqrt(sum(out[k]**2 for k in ('Q','K','V'))/3)
    if 'P' in base:out['JS']=js_rows(base['P'],changed['P']).transpose(0,2,1)
    return out


def link(lm,cache,layer,path,seed):
    started=time.monotonic();base=lm.recipient(cache['next_raw'],layer+1);hs=cache['heads'].shape[2];t=cache['raw'].shape[1];hr=base['Q'].shape[2]
    # Whole donor head across the full sequence is the single-ablation intervention.
    singles={k:[] for k in ('Q','K','V','QKV')+ (('JS',) if 'P' in base else ())}
    for start in range(0,hs,4):
        index=list(range(start,min(start+4,hs)));masks=np.zeros((len(index),t,hs),dtype=np.float32)
        for n,h in enumerate(index):masks[n,:,h]=1
        scores=distances(base,lm.intervene(cache,layer,masks))
        for k,v in scores.items():singles[k].extend(v)
    singles={k:np.asarray(v,dtype=np.float32) for k,v in singles.items()}
    # Identical cardinality weak/strong/random removals, ranking by balanced QKV influence.
    # One coherent full-sequence mask per recipient head; only that head's output scored.
    rank=singles['QKV'];rng=np.random.default_rng(seed)
    weak=masks_from_influences(rank,True);strong=masks_from_influences(rank,False)
    random=np.zeros_like(weak)
    for r in range(hr):
        for q in range(t):random[r,q,rng.choice(hs,hs//2,replace=False)]=1
    joint={};sum_individual={}
    for label,masks in [('weak_half',weak),('strong_half',strong),('random_half',random),('all_heads',np.ones_like(weak))]:
        values={k:[] for k in singles}
        for start in range(0,hr,4):
            stop=min(start+4,hr);scores=distances(base,lm.intervene(cache,layer,masks[start:stop]))
            for k,v in scores.items():
                for n,r in enumerate(range(start,stop)):values[k].append(v[n,:,r])
        for k,v in values.items():
            joint[label+'_'+k]=np.asarray(v,dtype=np.float32).T
            sum_individual[label+'_'+k]=np.sum(singles[k].transpose(2,1,0)*masks,axis=-1).T
    null=distances(base,lm.intervene(cache,layer,np.zeros((1,t,hs),dtype=np.float32)))
    if max(float(np.max(v)) for v in null.values())>2e-5:raise AssertionError('Null intervention changed recipient')
    payload={'single_'+k:v for k,v in singles.items()};payload.update({'joint_'+k:v for k,v in joint.items()});payload.update({'individual_sum_'+k:v for k,v in sum_individual.items()})
    payload.update({'mask_weak_half':weak,'mask_strong_half':strong,'mask_random_half':random})
    payload.update({'null_'+k:v for k,v in null.items()})
    np.savez_compressed(path,**payload)
    return {'source':layer,'recipient':layer+1,'source_heads':hs,'recipient_query_heads':hr,'full_attention_recipient':'P' in base,'elapsed_seconds':time.monotonic()-started,'null_max':max(float(np.max(v)) for v in null.values())}


def main():
    p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--out',type=Path,required=True);p.add_argument('--seconds',type=int,default=7200);a=p.parse_args()
    a.out.mkdir(parents=True,exist_ok=False);(a.out/'cache').mkdir();(a.out/'raw').mkdir();started=time.monotonic()
    protocol={'scope':'all 23 adjacent layer transitions; source attention head contribution to next effective Q/K/V','samples':'4 previously authored zh/en texts, first32 tokens, descriptive','primary_queries':list(range(7,32)), 'influence':'L2 change / baseline token recipient-head L2; QKV score sqrt(mean of three squared relative changes))','linear_QK':'after causal convolution/SiLU and delta-rule L2 normalization; Q additionally divided by sqrt(key dim)','full_QK':'QK norm and partial RoPE applied; K/V repeated to query heads for indexing (not independent KV observations)','single_ablation':'remove complete source head projected output at all tokens; recompute source postattention normalization, MLP, residual and next input normalization','joint_ablation':'weak/strong/random half and all attention heads; one coherent token-varying removal mask for each recipient query head; ranking uses QKV single effects; no fitting','continuous':'effective heads=1/sum(p^2); fraction=effective/H; concentration=(sum(p^2)-1/H)/(1-1/H)','zero_total':'undefined concentration (NaN arrays), count explicitly','cost_budget_seconds':a.seconds,'no_training':True,'keep_all_weights_and_caches':True,'limitations':['not final LM quality or pruning proof','individual effects not additive; normalized scores not true contribution probabilities','token-varying masks are input-specific diagnostic and not a deployable fixed head subset','linear attention does not have an ordinary softmax probability matrix','Q/K/V scales handled separately; combined QKV is a declared balanced relative-effect score']}
    dump(a.out/'protocol.json',protocol);status={'state':'running','completed_links':[]};dump(a.out/'status.json',status)
    try:
        lm=InfluenceLM(a.source/'weights',json.loads((a.source/'model_config.json').read_text()),seconds=a.seconds)
        samples=json.loads(Path('crossqkv/fresh_samples.json').read_text())['samples'];tok=Tokenizer.from_file(str(a.source/'tokenizer.json'));controls=[]
        for si,s in enumerate(samples):
            lm.guard();s['token_ids']=tok.encode(s['text']).ids[:32];ids=np.array([s['token_ids']],dtype=np.int64);cache=a.out/'cache'/s['id'];cache.mkdir()
            final=lm.baseline(ids,cache);np.save(cache/'final.npy',final)
            for x in lm.head_sum_errors:controls.append({'sample':s['id'],**x})
            if si==0:
                reference,_=StreamingLM.forward(lm,ids);normalized=lm.normalization(final,'norm');err=float(np.linalg.norm(normalized-reference)/max(np.linalg.norm(reference),1e-30));assert err<2e-5;status['canonical_forward_relative_error']=err
            dump(a.out/'head-sum-controls.json',controls);dump(a.out/'samples.json',samples)
            for layer in range(23):
                target=a.out/'raw'/f'{s["id"]}-link-{layer:02d}.npz'
                with np.load(cache/f'layer-{layer}.npz') as data:loaded={k:data[k] for k in data.files}
                info=link(lm,loaded,layer,target,31000+si*100+layer)
                del loaded
                status['completed_links'].append({'sample':s['id'],**info});status['elapsed_seconds']=time.monotonic()-started;dump(a.out/'status.json',status)
                print('completed',s['id'],layer,'seconds',round(info['elapsed_seconds'],2),flush=True)
        status.update(state='completed',elapsed_seconds=time.monotonic()-started,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss);dump(a.out/'status.json',status)
    except BaseException as e:
        status.update(state='failed',error=repr(e),elapsed_seconds=time.monotonic()-started);dump(a.out/'status.json',status);raise
if __name__=='__main__':main()