"""Paired CPU-only existing-adapter comparison for the published sqrt(JS) metric."""
import argparse,json,hashlib,time,resource,gc
from pathlib import Path
import numpy as np
from tokenizers import Tokenizer
from dominance.core import InfluenceLM,concentration
from dominance.run import distances
from .lora import install,ExplicitLoRA,BASE_REVISION
SOURCES=(2,6,10,14,18,22)

def dump(p,d):p.write_text(json.dumps(d,indent=2,allow_nan=False)+'\n')
def single(lm,cache,layer,path):
    base=lm.recipient(cache['next_raw'],layer+1);hs=cache['heads'].shape[2];t=cache['raw'].shape[1];assert hs==16
    scores={k:[] for k in ('Q','K','V','QKV','JS')}
    for start in range(0,hs,4):
        masks=np.zeros((4,t,hs),dtype=np.float32)
        for n in range(4):masks[n,:,start+n]=1
        effects=distances(base,lm.intervene(cache,layer,masks))
        for k,v in effects.items():scores[k].extend(v)
    null=distances(base,lm.intervene(cache,layer,np.zeros((1,t,hs),dtype=np.float32)))
    assert max(float(np.max(v)) for v in null.values())<2e-5
    payload={'single_'+k:np.asarray(v,dtype=np.float32) for k,v in scores.items()};payload.update({'null_'+k:v for k,v in null.items()});np.savez_compressed(path,**payload)

def main():
    p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--parent',type=Path,required=True);p.add_argument('--assets',type=Path,required=True);p.add_argument('--out',type=Path,required=True);a=p.parse_args();a.out.mkdir(parents=True,exist_ok=False);(a.out/'cache').mkdir();(a.out/'raw').mkdir()
    started=time.monotonic();meta=json.loads((a.assets/'adapter_meta.json').read_text());assert meta['base_revision']==BASE_REVISION
    adapter_sha=hashlib.sha256((a.assets/'adapter.safetensors').read_bytes()).hexdigest();assert adapter_sha=='93a3f43a4356f84da0d13ecc622fee13e41ba15d237cdaad6094f243899c62d6'
    tok=Tokenizer.from_file(str(a.source/'tokenizer.json'));legacy=json.loads((a.parent/'samples.json').read_text());gsm=[]
    for s in legacy:assert s['token_ids']==tok.encode(s['text']).ids[:32]
    for index,line in enumerate((a.assets/'gsm8k-test.jsonl').read_text().splitlines()):
        q=json.loads(line)['question'];ids=tok.encode(q).ids
        if len(ids)<32:continue
        gsm.append({'id':f'gsm8k-test-{index:04d}','row_index':index,'text':q,'question_sha256':hashlib.sha256(q.encode()).hexdigest(),'token_ids':ids[:32]})
        if len(gsm)==4:break
    assert len(gsm)==4
    cohorts={'legacy':legacy,'gsm8k':gsm};dump(a.out/'samples.json',cohorts);dump(a.out/'adapter_meta.json',meta);dump(a.out/'gsm8k-source.json',json.loads((a.assets/'gsm8k-source.json').read_text()))
    protocol={'primary':'same six full-attention receivers and original4x32 token cohort as published8.931822/16','adapter_repo':'Avifenesh/gdn-beta-gate-qwen3.5-2b-grpo','adapter_revision':'038d8be9bd0edd3883fa019651b02bdf65609468','adapter_file':'lora/seed0/adapter.safetensors','adapter_sha256':adapter_sha,'base_revision':BASE_REVISION,'scale':2,'rank':16,'arm':'pure LoRA','training':'publisher-reported GSM8K GRPO; we do no training','sample_choice':'four original texts; first4 canonical GSM8K test questions with >=32tokens, length selection only, first32tokens; question only','source_layers':SOURCES,'primary_queries':list(range(7,32)),'main_metric':'concentration(sqrt(single_JS)); same FP32 tensors and sequence-wide donor-head deletion as prior run','cohort_pairing':'each comparison uses same question/tokenIDs before and after; base legacy arrays reused with elementwise checks','runtime':'read-only FP32 effective W+2BA; no BF16 re-rounding or quantization','controls':'head sum, null mask, explicit low-rank versus effective-weight full forward relL2<1e-4','budget_seconds':7200,'keep_all_weights_and_caches':True,'limitations':['one existing training seed','4 texts per cohort, 32 tokens each','dependent query/head observations, not iid textual samples','no generation or accuracy measurement','GDN targets only; ordinary attention weights stay unchanged','no pruning or functional interchangeability claim']}
    dump(a.out/'protocol.json',protocol);status={'state':'running','links':[]};dump(a.out/'status.json',status);cfg=json.loads((a.source/'model_config.json').read_text());controls=[];references=[]
    try:
        for cohort,samples in cohorts.items():
            for role in ('base','finetuned'):
                reuse=cohort=='legacy' and role=='base';lm=None
                if not reuse:
                    lm=InfluenceLM(a.source/'weights',cfg,seconds=max(1,7200-(time.monotonic()-started)))
                    if role=='finetuned':pairs,scale=install(lm,a.assets,meta)
                for si,s in enumerate(samples):
                    if time.monotonic()-started>7200:raise TimeoutError('Paired CPU budget')
                    if not reuse:
                        cache=a.out/'cache'/f'{cohort}-{role}-{s["id"]}';cache.mkdir();ids=np.array([s['token_ids']],dtype=np.int64);final=lm.baseline(ids,cache);np.save(cache/'final.npy',final)
                        controls.extend({'cohort':cohort,'role':role,'sample':s['id'],**v} for v in lm.head_sum_errors)
                        if cohort=='legacy' and role=='finetuned' and si==0:
                            ref=ExplicitLoRA(a.source/'weights',cfg,seconds=max(1,7200-(time.monotonic()-started)));ref.pairs,ref.scale=pairs,scale;expected,_=ref.forward(ids);actual=lm.normalization(final,'norm');error=float(np.linalg.norm(actual.astype(np.float64)-expected)/np.linalg.norm(expected.astype(np.float64)));assert error<1e-4
                            status['explicit_lora_forward_relative_error']=error;del ref;gc.collect()
                    for layer in SOURCES:
                        out=a.out/'raw'/f'{cohort}-{role}-{s["id"]}-{layer:02d}.npz'
                        if reuse:
                            old=a.parent/'raw'/f'{s["id"]}-link-{layer:02d}.npz'
                            with np.load(old) as d:payload={prefix+k:d[prefix+k] for prefix in ('single_','null_') for k in ('Q','K','V','QKV','JS')}
                            np.savez_compressed(out,**payload)
                            with np.load(out) as d:
                                for k,v in payload.items():np.testing.assert_array_equal(d[k],v)
                            references.append({'source':str(old.resolve()),'sha256':hashlib.sha256(old.read_bytes()).hexdigest(),'derived':out.name})
                        else:
                            with np.load(cache/f'layer-{layer}.npz') as d:loaded={k:d[k] for k in d.files}
                            single(lm,loaded,layer,out);del loaded
                        status['links'].append({'cohort':cohort,'role':role,'sample':s['id'],'source':layer,'path':str(out.relative_to(a.out))});status['elapsed_seconds']=time.monotonic()-started;dump(a.out/'status.json',status);print('completed',cohort,role,s['id'],layer,flush=True)
                    dump(a.out/'head-sum-controls.json',controls);dump(a.out/'baseline-references.json',references)
                if lm is not None:del lm;gc.collect()
        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()