"""Recompute saved concentration, masks, scale and scope checks; preserve caches."""
import json,hashlib,csv,sys
from pathlib import Path
import numpy as np
from .core import concentration


def main(root,source):
    root=Path(root);source=Path(source);status=json.loads((root/'status.json').read_text());assert status['state']=='completed'
    samples=json.loads((root/'samples.json').read_text());assert len(samples)==4 and all(len(s['token_ids'])==32 for s in samples)
    expected={(s['id'],i) for s in samples for i in range(23)};actual={(x['sample'],x['source']) for x in status['completed_links']};assert actual==expected and len(status['completed_links'])==92
    cfg=json.loads((source/'model_config.json').read_text())['text_config'];obs=dict(np.load(root/'observations.npz'));lookup=sorted(s['id'] for s in samples);counts={};undefined={};full_js_links=0;joint_values={}
    for sample,layer in sorted(actual):
        f=root/'raw'/f'{sample}-link-{layer:02d}.npz';d=dict(np.load(f));hs=8 if cfg['layer_types'][layer]=='full_attention' else 16;hr=8 if cfg['layer_types'][layer+1]=='full_attention' else 16
        if hr==8:
            full_js_links+=1;d['single_sqrtJS']=np.sqrt(d['single_JS'])
            for label in ('weak_half','strong_half','random_half','all_heads'):
                d['joint_'+label+'_sqrtJS']=np.sqrt(d['joint_'+label+'_JS'])
                mask=d.get('mask_'+label,np.ones((hr,32,hs),dtype=np.float32))
                d['individual_sum_'+label+'_sqrtJS']=np.sum(d['single_sqrtJS'].transpose(2,1,0)*mask,-1).T
            d['null_sqrtJS']=np.sqrt(d['null_JS'])
        masks={label:d['mask_'+label] for label in ('weak_half','strong_half','random_half')}
        for label,m in masks.items():
            assert m.shape==(hr,32,hs);assert np.isin(m,[0,1]).all();np.testing.assert_array_equal(m.sum(-1),hs//2)
        np.testing.assert_array_equal(masks['weak_half']+masks['strong_half'],1)
        ranked=np.argsort(d['single_QKV'].transpose(2,1,0),axis=-1,kind='stable')
        np.testing.assert_array_equal(np.take_along_axis(masks['weak_half'],ranked[...,:hs//2],-1),1)
        for key in ('Q','K','V','QKV')+(('JS','sqrtJS') if hr==8 else ()):
            single=d['single_'+key];assert single.shape==(hs,32,hr) and np.isfinite(single).all() and np.min(single)>=0
            effect=single[:,7:,:].transpose(1,2,0);stats=concentration(effect);counts[key]=counts.get(key,0)+effect.shape[0]*hr;undefined[key]=undefined.get(key,0)+int((~stats['valid']).sum())
            choose=(obs[key+'_source']==layer)&(obs[key+'_sample_index']==lookup.index(sample))
            assert int(choose.sum())==25*hr
            for metric,values in stats.items():
                if metric=='valid':continue
                np.testing.assert_allclose(obs[key+'_'+metric][choose],values.ravel(),rtol=1e-6,atol=1e-7,equal_nan=True)
            for label in ('weak_half','strong_half','random_half','all_heads'):
                mask=masks.get(label,np.ones((hr,32,hs),dtype=np.float32));expected_sum=np.sum(single.transpose(2,1,0)*mask,-1).T
                np.testing.assert_allclose(d['individual_sum_'+label+'_'+key],expected_sum,rtol=1e-6,atol=1e-8)
                assert d['joint_'+label+'_'+key].shape==(32,hr) and np.isfinite(d['joint_'+label+'_'+key]).all()
                joint_values.setdefault((key,label),[]).extend(d['joint_'+label+'_'+key][7:].ravel())
            assert float(np.max(d['null_'+key]))<=2e-5
    assert full_js_links==24
    summary=json.loads((root/'analysis-summary.json').read_text())
    def check_stats(values,record):
        v=np.asarray(values,dtype=np.float64);v=v[np.isfinite(v)];assert len(v)==record['count']
        if len(v):
            for name,x in zip(['p10','median','p90','mean'],[*np.quantile(v,[.1,.5,.9]),v.mean()]):np.testing.assert_allclose(record[name],x,rtol=2e-6,atol=1e-8)
    for key in counts:
        for metric,record in summary['concentration'][key].items():check_stats(obs[key+'_'+metric],record)
    primary=json.loads((root/'primary-analysis-summary.json').read_text())
    for section in ('concentration','joint','paired_joint','by_recipient_type'):
        for key,value in primary[section].items():assert summary[section][key]==value,(section,key)
    full_sources=sorted({x['source'] for x in status['completed_links'] if x['full_attention_recipient']})
    assert summary['full_attention_source_indices']==full_sources
    for key,groups in summary['by_recipient_type'].items():
        src=obs[key+'_source']
        for label,metrics in groups.items():
            choose=np.isin(src,full_sources) if label=='full_attention' else ~np.isin(src,full_sources)
            for metric,record in metrics.items():check_stats(obs[key+'_'+metric][choose],record)
    for (key,label),values in joint_values.items():check_stats(values,summary['joint'][key][label])
    for key in counts:
        weak=np.asarray(joint_values[(key,'weak_half')]);strong=np.asarray(joint_values[(key,'strong_half')]);random=np.asarray(joint_values[(key,'random_half')]);valid=weak>0;record=summary['paired_joint'][key]
        assert record['zero_weak_effect_count']==int((~valid).sum())
        np.testing.assert_allclose(record['weak_less_than_strong_fraction'],np.mean(weak<strong))
        np.testing.assert_allclose(record['weak_less_than_random_fraction'],np.mean(weak<random))
        check_stats(strong[valid]/weak[valid],record['strong_over_weak']);check_stats(random[valid]/weak[valid],record['random_over_weak'])
    controls=json.loads((root/'head-sum-controls.json').read_text());assert len(controls)==96 and max(x['relative_l2'] for x in controls)<2e-5
    assert status['canonical_forward_relative_error']<2e-5
    caches=[]
    for f in sorted((root/'cache').rglob('*')):
        if f.is_file():caches.append({'path':str(f.resolve()),'bytes':f.stat().st_size,'sha256':hashlib.sha256(f.read_bytes()).hexdigest()})
    weights=list((source/'weights').rglob('*.safetensors'));assert weights and sum(f.stat().st_size for f in weights)==4548221488
    manifest=json.loads((source/'manifest.json').read_text());weight_records=[]
    for f in weights:
        with f.open('rb') as stream:sha=hashlib.file_digest(stream,'sha256').hexdigest()
        expected=next(x['upstream_sha256'] for x in manifest['files'] if x['name']==f.name);assert sha==expected
        weight_records.append({'path':str(f.resolve()),'bytes':f.stat().st_size,'sha256':sha,'matches_upstream':True})
    external={'host':'8.211.154.134','checkpoint_directory':str((source/'weights').resolve()),'checkpoint_bytes':sum(f.stat().st_size for f in weights),'checkpoint_files':weight_records,'revision':'15852e8c16360a2fea060d615a32b45270f8a8fc','baseline_caches':caches,'automatic_deletion':False,'prior_caches_untouched':True}
    (root/'external-caches.json').write_text(json.dumps(external,indent=2)+'\n')
    reference=json.loads((source/'reference-check.json').read_text())
    assert reference['passed'] and all(x['passed'] for x in reference['tests'])
    assert hashlib.sha256(Path('functional/model.py').read_bytes()).hexdigest()==reference['implementation_sha256']
    (root/'prior-official-reference-check.json').write_bytes((source/'reference-check.json').read_bytes())
    files=list(Path('dominance').glob('*.py'))+[Path('dominance/PROTOCOL.md'),Path('dominance/README.md'),Path('tests/test_dominance.py'),Path('functional/model.py'),Path('crossqkv/fresh_samples.json')]
    (root/'source-hashes.json').write_text(json.dumps({str(f):hashlib.sha256(f.read_bytes()).hexdigest() for f in files if not f.name.startswith('._')},indent=2)+'\n')
    audit={'passed':True,'transitions':23,'texts':4,'link_observations':92,'primary_observations_per_metric':counts,'undefined_zero_total_per_metric':undefined,'full_attention_js_links':full_js_links,'head_sum_controls':96,'canonical_forward_relative_error':status['canonical_forward_relative_error'],'all_raw_concentration_and_masks_recomputed':True,'pooled_concentration_and_joint_summary_recomputed':True,'original_primary_summaries_unchanged':True,'sqrt_JS_scale_sensitivity_recomputed':True,'retained_baseline_cache_files':len(caches),'checkpoint_sha_matches_upstream':True,'prior_official_reference_check_passed_same_implementation':True,'prior_official_reference_tests':len(reference['tests'])}
    (root/'audit.json').write_text(json.dumps(audit,indent=2)+'\n');print(json.dumps(audit,indent=2),flush=True)
    included=[f for f in root.rglob('*') if f.is_file() and 'cache' not in f.relative_to(root).parts and f.name!='checksums.json']
    (root/'checksums.json').write_text(json.dumps({str(f.relative_to(root)):hashlib.sha256(f.read_bytes()).hexdigest() for f in sorted(included)},indent=2)+'\n')
if __name__=='__main__':main(sys.argv[1],sys.argv[2])