"""Explicit adapters for the released Lab17 core and corrected Lab19 schema.

This narrow adapter accepts only reviewed source identities, not a purported
new experiment with coincidentally compatible shapes. Imported Python is never
executed. Missing original file hashes are explicitly recorded at import.
"""
from pathlib import Path
from microscope_io import bounded,digest,read_json,relative,npz,events,encode,PAYLOAD_CAP

REPOSITORY='EleutherAI/pythia-14m'
REVISION='94f7c35d5e9f2e9bac8ca839329f505b4d007d5d'
ARTIFACTS={
 'config.json':{'bytes':698,'sha256':'f97f966a66c444890ed461fff2a51eefb15d74303df05b948124719f199b0b17'},
 'generation_config.json':{'bytes':111,'sha256':'494cbba887cc299273f0c54d7612596918e74c7881a27cbd95e4e7e2464ea090'},
 'special_tokens_map.json':{'bytes':441,'sha256':'10b8c8852c1e1f70b54d9aff61728408c28971c0e97a6c5a7b2debbd1d3e9c0c'},
 'tokenizer.json':{'bytes':2114042,'sha256':'870f4e2baa6b683221fa52004d5d6f40ab8c9d31961617304b78c910c2c3caf2'},
 'tokenizer_config.json':{'bytes':4834,'sha256':'eee017c5bd133137f45907bd0a6e781e2ccd1a533734b7ed2a2f2f4446659809'},
 'model.safetensors':{'bytes':28143920,'sha256':'116a02532db461f91386a5b20f942ff2c8d4de7341e21b55caafc3d7b25f49a1'}}
VERSIONS={'torch':'2.9.0','transformers':'4.57.1','tokenizers':'0.22.1','safetensors':'0.6.2','huggingface-hub':'0.35.3'}
SOURCE17={'inspect_activations.py':'3c92dd075e8d017b9253cb0e6c44930feb20a45fe9a7ad443de54c69322c3ab1',
 'inspection_common.py':'792f5dc299fc018bb4b91b933a6bddff78ddc6d27e59e71be950c13df854ac56',
 'inspect_residual_stream.py':'4ee3f277831116c7b27cd0a66ef7dc80452f4659116cbcda36d9e52750d04ff5'}
SOURCE19={'contrast_experiment.py':'10c1531bce1e0ae324275a8661852205be8a3d06188f2eb2699ab7d61f0bdc32',
 'inspection_common.py':'c608aa7d57efa31ca6ae9a3825cf401c7c3336ae4170a658dee33925e3b0f30c',
 'inspect_residual_stream.py':SOURCE17['inspect_residual_stream.py'],
 'run_lab19.py':'af3a6e878694a16cb8ec938e86a70992c38491e2e01e75fffba3850da2fd9f9d',
 'lab19-fixtures.json':'d15f961a7bb80b27a15059bd5fa0345d7d135f0c0c19c2539ad097bc15e30d58'}
FILES17=('captures.json','checks.json','manifest.json','outcome.json','prompts.json','predictions.json')
FILES19=('artifact-manifest.json','environment.json','fixtures.json','tokens.json','direction.json','sources.json',
 'discovery-measurements.json','discovery-vectors.npz','discovery-pilot.json','predictions.json','predictions.json.sha256',
 'plan-draft.json','state.json','heldout-start.json','attempts.jsonl','array-index.json','logits.npz','vectors.npz',
 'measurements.json','checks.json','metrics.csv','summary.json')
CAPTURE_KEYS=tuple([f'r{i}' for i in range(7)]+[f'a{i}' for i in range(6)]+[f'm{i}' for i in range(6)]+['h_f'])
ARMS19=('baseline','zero','positive_small','negative_small','positive_larger','negative_larger','random_positive','random_negative','remove')
DISCOVERY=tuple(f'D{i}{sign}' for i in range(1,4) for sign in ('+','-'))
HELDOUT=('H1+','H1-','H2+','H2-','O1','O2')
DEMO_TEXTS={'H1+':'The concert was wonderful and the musicians were talented.\nOverall, it was',
 'H1-':'The concert was awful and the musicians were unskilled.\nOverall, it was'}

def original_schedule():
    return ([(p,'baseline') for p in DISCOVERY]+[(p,'observe') for p in DISCOVERY]+
        [('D1+','baseline'),('D1+','zero'),('D1+','abort'),('D1+','baseline')]+
        [(p,'positive_small') for p in DISCOVERY]+[(p,a) for p in HELDOUT for a in ARMS19]+
        [(p,'baseline') for p in HELDOUT])

def tokens(row,max_tokens=64):
    ids=row.get('token_ids');n=len(ids) if isinstance(ids,list) else 0
    if not 1<=n<=max_tokens or any(type(i) is not int or not 0<=i<50304 for i in ids):raise ValueError('Token inventory/budget')
    if row.get('positions')!=list(range(n)) or len(row.get('decoded_pieces',[]))!=n:raise ValueError('Token pieces/positions')
    derived={}
    for key,value in [('length',n),('final_position',n-1),('attention_mask',[1]*n)]:
        if key in row and row[key]!=value:raise ValueError('Saved token field disagrees: '+key)
        if key not in row:derived[key]='derived from saved token_ids (no padding)'
    return {**row,'length':n,'final_position':n-1,'attention_mask':[1]*n,'derived_fields':derived,
            'prompt_provenance':'authored_synthetic'}

def specimen(manifest):
    if manifest.get('repository')!=REPOSITORY or manifest.get('revision')!=REVISION or manifest.get('artifacts')!=ARTIFACTS:
        raise ValueError('Imported model identity differs')

def compatible(env):
    if env.get('python')!='3.13.13' or env.get('packages')!=VERSIONS:raise ValueError('Historical runtime differs')
    for k,v in {'device':'cpu','dtype':'float32','intra_op':1,'inter_op':1}.items():
        if env.get(k)!=v:raise ValueError('Historical computational setting differs: '+k)

def address(key):
    if key not in CAPTURE_KEYS:raise ValueError('Unknown capture boundary')
    if key=='h_f':return 'gpt_neox.final_layer_norm',None,'normalized'
    i=int(key[1:]);kind=key[0]
    if kind=='r':return ('gpt_neox.layers.0/input',0,'raw') if i==0 else (f'gpt_neox.layers.{i-1}/output',i-1,'raw')
    return f'gpt_neox.layers.{i}.post_'+('attention' if kind=='a' else 'mlp')+'_dropout',i,'raw_branch_contribution'

def passed(record):
    return isinstance(record,dict) and record.get('passed') is True and isinstance(record.get('max_abs_error'),(int,float)) and record['max_abs_error']>=0

def lab17(root):
    import numpy as np
    from inspect_activations import CORE
    get=lambda n:read_json(relative(root,n))
    manifest=get('manifest.json');specimen(manifest);compatible(manifest)
    if manifest.get('sources')!=SOURCE17 or manifest.get('seed')!=0 or manifest.get('attn_backend')!='eager' or manifest.get('use_cache') is not False:
        raise ValueError('Unsupported Lab17 source/settings')
    outcome=get('outcome.json')
    if outcome.get('status')!='completed' or outcome.get('core_prompts_completed')!=8 or outcome.get('failure') is not None:
        raise ValueError('Lab17 core prerequisite incomplete')
    prompts=get('prompts.json')['core']
    if [(p['id'],p['text']) for p in prompts]!=[(f'P{i+1}',s) for i,s in enumerate(CORE)]:raise ValueError('Lab17 exact core prompts')
    prompts=[tokens(p,96) for p in prompts];captures=get('captures.json');checks=get('checks.json')
    if set(captures)!={p['id'] for p in prompts}:raise ValueError('Lab17 prompt coverage')
    vectors={}
    for row in prompts:
        p=row['id'];values=captures[p]
        if set(values)!=set(CAPTURE_KEYS):raise ValueError('Lab17 capture inventory')
        for key,record in values.items():
            module,block,_=address(key)
            required={'module_path':module,'boundary':key,'block_index':block,'prompt':p,'token_id':row['token_ids'][-1],
                      'position':row['final_position'],'shape':[128],'source_shape':[1,row['length'],128],'dtype':'float32'}
            if any(record.get(k)!=v for k,v in required.items()):raise ValueError('Ambiguous Lab17 capture address: '+p+'/'+key)
            vector=np.asarray(record.get('vector'),dtype=np.float32)
            if vector.shape!=(128,) or not np.isfinite(vector).all():raise ValueError('Invalid Lab17 vector')
            vectors[p+'_'+key]=vector
        matches=[c for c in checks if c.get('prompt')==p]
        if len(matches)!=1:raise ValueError('Lab17 check coverage')
        c=matches[0]
        if c.get('hooks_removed') is not True or not all(passed(c.get(k)) for k in ('hooked_vs_baseline','after_removal_vs_baseline','final_normalization','normalized_logits','raw_logits')):
            raise ValueError('Lab17 invariance/reconstruction prerequisite failed')
        if set(c.get('addition',{}))!=set(map(str,range(6))) or not all(passed(x) for x in c['addition'].values()):raise ValueError('Lab17 addition checks')
        for i in range(6):
            if not np.allclose((vectors[p+f'_m{i}']+vectors[p+f'_a{i}'])+vectors[p+f'_r{i}'],vectors[p+f'_r{i+1}'],atol=1e-6,rtol=1e-5):
                raise ValueError('Imported Lab17 residual identity')
    # Original full logit arrays were not saved; reconstruction checks remain
    # declared historical evidence, not newly recomputed logit comparisons.
    return {'original_status':outcome['status'],'declared_status_file':'outcome.json','completeness':'derived: eight core prompts and required checks',
        'origin_run_id':'lab17-'+digest(bounded(relative(root,'manifest.json'))),
        'provenance':'measured','predictions_frozen_record':'predictions.json','optional_probe':'excluded',
        'historical_sources':manifest['sources'],'historical_numpy':'not recorded',
        'compatibility_notes':['Old platform.platform/processor metadata replaced by os.uname metadata; numerical helpers unchanged.',
          'Current inspect_activations optional-probe persistence differs; core capture/run_forward/checked_logits are unchanged.',
          'Historical seed0; fresh integration seed19 is declared. Eval has no stochastic sampling.'],
        'prompts':prompts,'vectors':vectors,'file_hashes':{n:{'sha256':digest(bounded(relative(root,n))),
            'integrity':'recorded_at_import'} for n in FILES17}}

def lab19(root):
    import numpy as np
    get=lambda n:read_json(relative(root,n))
    specimen(get('artifact-manifest.json'));compatible(get('environment.json'))
    sources=get('sources.json')
    if sources!=SOURCE19:raise ValueError('Unsupported Lab19 source identity; corrected runner required')
    state=get('state.json');predictions=get('predictions.json');prediction_hash=digest(bounded(relative(root,'predictions.json')))
    if state.get('lab')!=19 or state.get('phase')!='completed' or state.get('attempts')!=82:raise ValueError('Lab19 completed core prerequisite')
    if bounded(relative(root,'predictions.json.sha256'),1000).decode().strip()!=prediction_hash or state.get('plan_sha256')!=prediction_hash:
        raise ValueError('Original frozen prediction hash mismatch')
    if predictions.get('sources')!=sources or predictions.get('block')!=2 or predictions.get('boundary')!='raw post-block residual after additions before block 3' or predictions.get('score_ids')!=[1175,3076]:
        raise ValueError('Lab19 frozen computational address')
    if predictions.get('atol')!=1e-6 or predictions.get('rtol')!=1e-6 or predictions.get('attempts_at_freeze')!=22:
        raise ValueError('Lab19 frozen schedule/tolerances')
    start=get('heldout-start.json')
    if start.get('plan_sha256')!=prediction_hash or start.get('attempt_count_before')!=22 or start['timestamp']<=predictions['timestamp']:
        raise ValueError('Original evaluation/freeze chronology')
    if digest(bounded(relative(root,'plan-draft.json')))!=state.get('plan_draft_sha256'):raise ValueError('Lab19 plan-draft hash')
    for n,h in predictions['files'].items():
        if n not in FILES19:raise ValueError('Unapproved original manifest file')
        if digest(bounded(relative(root,n)))!=h:raise ValueError('Changed imported bytes: '+n)
    records=events(relative(root,'attempts.jsonl'));schedule=original_schedule()
    if len(records)!=164:raise ValueError('Lab19 chronological attempts incomplete')
    completed=[]
    for number,(prompt,arm) in enumerate(schedule,1):
        a,b=records[2*(number-1):2*number]
        expected_status='expected_exception' if number==15 else 'completed'
        if any(e.get('attempt')!=number or e.get('prompt')!=prompt or e.get('arm')!=arm for e in (a,b)) or a.get('status')!='started' or b.get('status')!=expected_status:
            raise ValueError('Lab19 fixed chronological schedule')
        if b.get('hook_cleanup') is not True or b.get('hook_calls')!=(0 if arm=='baseline' else 1):raise ValueError('Lab19 hook checks')
        if b['timestamp']<a['timestamp'] or (number>1 and a['timestamp']<records[2*number-3]['timestamp']):raise ValueError('Lab19 event chronology')
        if number==15:
            if b.get('exception')!='deliberate_cleanup_check' or b.get('array_file') is not None:raise ValueError('Lab19 cleanup exception')
        else:completed.append(b)
    index=get('array-index.json')
    if len(index)!=81 or any(any(record.get(k)!=event.get(k) for k in ('attempt','prompt','arm','status','hook_cleanup','hook_calls')) or record.get('key')!=f"attempt_{event['attempt']:03d}" for record,event in zip(index,completed)):
        raise ValueError('Lab19 array-index provenance')
    logit_schema={f'attempt_{n:03d}':((50304,),'float32') for n in range(1,83) if n!=15}
    logits,payload=npz(relative(root,'logits.npz'),logit_schema)
    vector_schema={}
    for n,(_,arm) in enumerate(schedule,1):
        if arm not in ('baseline','abort'):
            vector_schema[f'attempt_{n:03d}_original']=((128,),'float32')
            if arm!='observe':vector_schema[f'attempt_{n:03d}_replacement']=((128,),'float32')
    vectors,count=npz(relative(root,'vectors.npz'),vector_schema);payload+=count
    discovery,count=npz(relative(root,'discovery-vectors.npz'),{p:((128,),'float32') for p in DISCOVERY});payload+=count
    if payload>PAYLOAD_CAP:raise ValueError('Imported numeric payload cap')
    d=np.mean([discovery[f'D{i}+'].astype(np.float64)-discovery[f'D{i}-'].astype(np.float64) for i in range(1,4)],axis=0)
    norm=float(np.linalg.norm(d));R=float(np.mean([np.linalg.norm(discovery[p].astype(np.float64)) for p in DISCOVERY]))
    if norm<=1e-8 or R<=1e-8:raise ValueError('Degenerate imported direction')
    direction=get('direction.json');derived={'d':d,'u':d/norm,'R':R,'dose':.05*R,'norm':norm}
    for k,v in derived.items():
        saved=np.asarray(direction.get(k),dtype=np.float64)
        if saved.shape!=np.asarray(v).shape or not np.isfinite(saved).all() or not np.allclose(v,saved,atol=1e-10,rtol=1e-10):raise ValueError('Imported direction/dose mismatch: '+k)
    if predictions['direction']!=direction:raise ValueError('Direction changed after freeze')
    rows=get('tokens.json')['rows']
    if [r['id'] for r in rows]!=list(DISCOVERY)+list(HELDOUT):raise ValueError('Lab19 complete fixture inventory')
    rows=[tokens(r) for r in rows];demo=[r for r in rows if r['id'] in DEMO_TEXTS]
    if [(r['id'],r['text']) for r in demo]!=list(DEMO_TEXTS.items()) or len({r['token_ids'][-1] for r in demo})!=1:
        raise ValueError('Lab19 exact demonstration inputs')
    fixtures=get('fixtures.json')
    if digest(encode(fixtures))!=digest(encode(read_json(Path(__file__).with_name('lab19-fixtures.json')))):raise ValueError('Lab19 synthetic fixtures differ')
    measurements=get('measurements.json')
    if [(m.get('prompt'),m.get('arm')) for m in measurements]!=[(p,a) for p in HELDOUT for a in ARMS19]:raise ValueError('Lab19 all review/control arms required')
    checks=get('checks.json');expected_names=[p+' observer' for p in DISCOVERY]+['discovery repeat','discovery zero','discovery recovery']+[p+' zero' for p in HELDOUT]+[p+' restoration' for p in HELDOUT]
    if [c.get('name') for c in checks]!=expected_names or not all(passed(c) for c in checks):raise ValueError('Lab19 numerical checks incomplete')
    return {'original_status':state['phase'],'declared_status_file':'state.json','standalone_status':'missing',
        'completeness':'derived: exact82 attempts,81logits,all54heldout arms and21checks',
        'origin_run_id':'lab19-'+prediction_hash,'provenance':'measured','predictions_sha256':prediction_hash,
        'historical_sources':sources,'historical_numpy':'not recorded','prior_exposure':'POST-EXPOSURE VALIDATION: same inspected fixtures; not fresh blinded evidence',
        'attention_mask_convention':'explicit saved all-ones mask passed by the reviewed corrected runner',
        'demo':demo,'rows':rows,'logits':logits,'vectors':vectors,'discovery':discovery,'direction':{k:np.asarray(direction[k],dtype=np.float64) if k in ('d','u') else direction[k] for k in ('d','u','R','dose','norm')},
        'attempts':completed,'payload_bytes':payload,'direction_verification':{k:np.asarray(v).tolist() for k,v in derived.items()},
        'file_hashes':{n:{'sha256':digest(bounded(relative(root,n))),
          'integrity':'matched_original_manifest' if n in predictions['files'] or n in ('predictions.json','plan-draft.json') else 'recorded_at_import'} for n in FILES19}}

def snapshot(store,root17,root19):
    sources=[(root17,'lab17',FILES17),(root19,'lab19',FILES19)]
    selected=[]
    for root,kind,names in sources:
        for n in names:
            source=relative(root,n);data=bounded(source);selected.append((f'imports/{kind}/{n}',data))
    if sum(len(data) for _,data in selected)+1_000_000>store.cap:raise ValueError('Over-budget approved import snapshot')
    # Validate the originals before copying; validate frozen copies again later.
    a,b=lab17(root17),lab19(root19)
    for name,data in selected:store.put(name,data)
    for result in (a,b):
        for key in ('vectors','logits','discovery','direction'):result.pop(key,None)
    store.json('imports.json',{'schema':'llm-microscope/1','lab17':a,'lab19':b,
        'mapping':{'lab17':'captures.json:P1..P8, twenty addressed final-position vectors; optional probe excluded',
          'lab19':'array-index.json chronological attempt_NNN -> logits.npz; vectors original/replacement; discovery D1..D3 +/-'},
        'historical_predicate':'Reviewed source schemas and actual recorded evidence; hashes do not authenticate scientific authorship'})
    return a,b