"""Independent paired effective-head, source, preservation and delivery checks."""
import json,hashlib,sys
from pathlib import Path
import numpy as np

def sha(path):
    with Path(path).open('rb') as f:return hashlib.file_digest(f,'sha256').hexdigest()
def main(root,parent,assets,source):
    root,parent,assets,source=map(Path,[root,parent,assets,source]);status=json.loads((root/'status.json').read_text());assert status['state']=='completed' and len(status['links'])==96
    samples=json.loads((root/'samples.json').read_text());assert set(samples)=={'legacy','gsm8k'} and all(len(v)==4 for v in samples.values());assert all(len(s['token_ids'])==32 for ss in samples.values() for s in ss)
    assert status['explicit_lora_forward_relative_error']<1e-4;rawstats={};expected=set();obs=dict(np.load(root/'observations.npz'));summary=json.loads((root/'summary.json').read_text())
    def check(x,record):
        x=np.asarray(x,dtype=np.float64);x=x[np.isfinite(x)];assert record['count']==len(x)
        for name,value in zip(('p10','median','p90','mean'),[*np.quantile(x,[.1,.5,.9]),x.mean()]):np.testing.assert_allclose(record[name],value,atol=1e-8,rtol=1e-6)
    for cohort,ss in samples.items():
        for role in ('base','finetuned'):
            for sample in ss:
                for layer in (2,6,10,14,18,22):
                    name=f'{cohort}-{role}-{sample["id"]}-{layer:02d}.npz';expected.add(name);d=dict(np.load(root/'raw'/name))
                    for key in ('Q','K','V','QKV','JS'):
                        assert d['single_'+key].shape==(16,32,8) and np.isfinite(d['single_'+key]).all() and np.min(d['single_'+key])>=0
                        assert np.max(d['null_'+key])<2e-5
                    if cohort=='legacy' and role=='base':
                        with np.load(parent/'raw'/f'{sample["id"]}-link-{layer:02d}.npz') as original:
                            for key in d:np.testing.assert_array_equal(d[key],original[key])
                    for metric,impacts in [('distance',np.sqrt(d['single_JS'])),('divergence',d['single_JS']),('QKV',d['single_QKV'])]:
                        x=impacts[:,7:,:].transpose(1,2,0).astype(np.float64);total=x.sum(-1);den=np.sum(x*x,-1);assert np.all(den>0)
                        neff=total*total/den
                        rawstats.setdefault((cohort,role,metric),[]).append(neff.ravel())
    assert expected=={f.name for f in (root/'raw').glob('*.npz')}
    for cohort in samples:
        for metric in ('distance','divergence','QKV'):
            for role in ('base','finetuned'):
                x=np.concatenate(rawstats[(cohort,role,metric)]);assert len(x)==4800
                np.testing.assert_allclose(obs['__'.join((cohort,role,metric,'effective_heads'))],x,rtol=1e-8,atol=1e-8)
                check(x,summary[cohort][metric][role]['effective_heads'])
            before=np.concatenate(rawstats[(cohort,'base',metric)]);after=np.concatenate(rawstats[(cohort,'finetuned',metric)]);check(after-before,summary[cohort][metric]['paired']['effective_heads_delta'])
    base=summary['legacy']['distance']['base']['effective_heads'];np.testing.assert_allclose([base[x] for x in ('median','p10','p90')],[8.93182211792415,5.7840811612328125,11.67265834734814],rtol=1e-10)
    controls=json.loads((root/'head-sum-controls.json').read_text());assert len(controls)==288 and all(c['relative_l2']<2e-5 for c in controls)
    refs=json.loads((root/'baseline-references.json').read_text());assert len(refs)==24 and all(sha(v['source'])==v['sha256'] for v in refs)
    assert sha(assets/'adapter.safetensors')=='93a3f43a4356f84da0d13ecc622fee13e41ba15d237cdaad6094f243899c62d6'
    dataset=json.loads((root/'gsm8k-source.json').read_text());assert sha(assets/'gsm8k-test.jsonl')==dataset['sha256']
    oldweights=list((source/'weights').glob('*.safetensors'));assert len(oldweights)==1
    base_sha=sha(oldweights[0]);assert base_sha=='aa33250c4fc64891ddfaba3a314fd9542ea371843c387178b425fbcc5ed680b1'
    from pilot.tensors import load_tensors
    import re
    tensors=load_tensors(assets);pairs={}
    for name,t in tensors.items():
        stem,which=name.split('.lora_');pairs.setdefault(stem,{})[which[0]]=t.read().reshape(t.shape).copy();t.evict()
    updates=[]
    for name,pair in pairs.items():
        a=pair['A'].astype(np.float64);b=pair['B'].astype(np.float64)
        norm2=4*float(np.sum((b.T@b)*(a@a.T).T))
        updates.append({'module':name,'delta_frobenius':float(np.sqrt(max(norm2,0))),'A_nonzero':int(np.count_nonzero(a)),'B_nonzero':int(np.count_nonzero(b))})
    assert len(updates)==72 and any(x['delta_frobenius']>0 for x in updates)
    (root/'weight-update-summary.json').write_text(json.dumps({'scale':2,'modules':updates,'nonzero_update_modules':sum(x['delta_frobenius']>0 for x in updates),'norm_identity':'||2BA||_F^2 = 4 trace((B^T B)(A A^T))'},indent=2)+'\n')
    changes=[]
    for cohort,ss in samples.items():
        for sample in ss:
            basefile=parent/'cache'/sample['id']/'final.npy' if cohort=='legacy' else root/'cache'/f"{cohort}-base-{sample['id']}"/'final.npy'
            ftfile=root/'cache'/f"{cohort}-finetuned-{sample['id']}"/'final.npy'
            before=np.load(basefile).astype(np.float64);after=np.load(ftfile).astype(np.float64)
            changes.append({'cohort':cohort,'sample':sample['id'],'final_pre_norm_residual_relative_l2':float(np.linalg.norm(after-before)/np.linalg.norm(before))})
    assert any(x['final_pre_norm_residual_relative_l2']>0 for x in changes)
    (root/'hidden-state-changes.json').write_text(json.dumps({'definition':'final residual before final RMSNorm, not logits, accuracy or LM quality','changes':changes},indent=2)+'\n')
    caches=[{'path':str(p.resolve()),'bytes':p.stat().st_size,'sha256':sha(p)} for p in sorted((root/'cache').rglob('*')) if p.is_file()]
    assert len(caches)==300
    (root/'external-caches.json').write_text(json.dumps({'host':'8.211.154.134','base_weights':str(oldweights[0].resolve()),'base_sha256':base_sha,'adapter_path':str((assets/'adapter.safetensors').resolve()),'adapter_sha256':sha(assets/'adapter.safetensors'),'caches':caches,'automatic_deletion':False},indent=2)+'\n')
    files=list(Path('finetune_compare').glob('*.py'))+[Path('finetune_compare/PROTOCOL.md'),Path('tests/test_finetune_compare.py'),Path('dominance/core.py'),Path('dominance/run.py'),Path('functional/model.py')]
    (root/'source-hashes.json').write_text(json.dumps({str(f):sha(f) for f in files},indent=2)+'\n')
    proof={'passed':True,'links':96,'paired_observations_per_cohort_metric':4800,'legacy_baseline_exact':True,'effective_head_formula_independently_recomputed':True,'head_sum_checks':288,'explicit_lora_forward_relative_error':status['explicit_lora_forward_relative_error'],'adapter_sha_verified':True,'base_sha_verified':True,'dataset_sha_verified':True,'retained_cache_files':len(caches),'nonzero_update_modules':sum(x['delta_frobenius']>0 for x in updates),'nonzero_hidden_state_changes_verified':True}
    (root/'audit.json').write_text(json.dumps(proof,indent=2)+'\n');print(json.dumps(proof,indent=2),flush=True)
    included=[p for p in root.rglob('*') if p.is_file() and 'cache' not in p.relative_to(root).parts and p.name!='checksums.json'];(root/'checksums.json').write_text(json.dumps({str(p.relative_to(root)):sha(p) for p in sorted(included)},indent=2)+'\n')
if __name__=='__main__':main(*sys.argv[1:])