"""Preregister exact legacy APL data and unmodified fixed CAL requests."""
from datetime import datetime,timezone
import json
import math
import os
from pathlib import Path
import shutil
import subprocess
import sys
import numpy as np
import pandas as pd
import yaml
from context import ROOT,MODEL,EXP,AF_MODEL,AF_EXP,AD_EXP,CFG,save,digest,sources,integrity,budget

def arr(path,value):
    with Path(path).open('xb') as f:np.save(f,value,allow_pickle=False)

def graph(folder,nodes,full):
    folder.mkdir();nodes=np.asarray(nodes,np.int32);ptr,dst,src,raw,base=full
    if len(nodes)==CFG['neurons']:ids=np.arange(len(dst),dtype=np.int64)
    else:
        ids=np.concatenate([np.arange(ptr[s],ptr[s+1],dtype=np.int64) for s in nodes])
        ids=ids[np.isin(dst[ids],nodes)]
    lookup=np.full(CFG['neurons'],-1,np.int32);lookup[nodes]=np.arange(len(nodes),dtype=np.int32)
    ss,dd=lookup[src[ids]],lookup[dst[ids]];pp=np.r_[np.int64(0),np.bincount(ss,minlength=len(nodes)).cumsum(dtype=np.int64)]
    for name,value in dict(nodes=nodes,edge_id=ids,src=ss,dst=dd,ptr=pp,signed_raw=raw[ids],weights=base[ids]).items():arr(folder/(name+'.npy'),value)
    pd.DataFrame({'local_index':np.arange(len(nodes),dtype=np.int32)},index=pd.Index(nodes,name='global_neuron_idx')).to_csv(folder/'completeness.csv')
    pd.DataFrame({'Presynaptic_Index':ss,'Postsynaptic_Index':dd,'Excitatory x Connectivity':raw[ids],'edge_id':ids},copy=False).to_parquet(folder/'connectivity.parquet',index=False)
    save(folder/'audit.json',dict(nodes=len(nodes),edges=len(ids),canonical_original_edge_ids=True,anatomical_weights_not_modified=True,
        effective_APL_zeroing='Only in new runtime clone, never these original weight files'))
    return nodes

def case(name,kind,steps,nodes,graphdir,requests,initial=None,extra=None):
    folder=EXP/'cases'/name;folder.mkdir()
    model=json.loads((AD_EXP/'protocol/preregistered.json').read_text())['config']['model']
    spec=dict(name=name,kind=kind,steps=steps,dt_ms=.1,duration_ms=steps*.1,model=model,
        graph_dir=str(graphdir),apl_bundle_dir=str(EXP/'apl-bundle'),apl=CFG['apl'],record_history=kind=='fixture',
        coefficients_path=str(AF_EXP/'oracle/no_input/coefficients.json'),
        science_python=str(ROOT/'models/male-cns-lif/.venv/bin/python'),legacy_apl_worker=str(MODEL/'legacy_apl_worker.py'),
        budget_root=str(EXP),output_budget_bytes=CFG['output_budget_bytes'],precision='float64',apl_precision='float32',training=False)
    if extra:spec.update(extra)
    requests=np.asarray(requests,dtype=np.int32).reshape(-1,2)
    if len(requests):
        requests=requests[np.lexsort((requests[:,1],requests[:,0]))]
        assert len(np.unique(requests.astype(np.int64)[:,0]*CFG['neurons']+requests[:,1]))==len(requests)
        assert requests[:,0].min()>=0 and requests[:,0].max()<steps and np.isin(requests[:,1],nodes).all()
    with (folder/'requests.npz').open('xb') as f:np.savez_compressed(f,tick=requests[:,0],neuron_idx=requests[:,1])
    if kind=='fixture':
        v=np.full(len(nodes),model['resting_potential_mv'],np.float64);g=np.zeros(len(nodes),np.float64)
        if initial:
            for node,value in initial.get('v',{}).items():v[np.flatnonzero(nodes==node)[0]]=value
            for node,value in initial.get('g',{}).items():g[np.flatnonzero(nodes==node)[0]]=value
        with (folder/'arrays.npz').open('xb') as f:np.savez_compressed(f,initial_v_mv=v,initial_g_mv=g)
    save(folder/'spec.json',spec)

def main():
    assert not (EXP/'protocol/preregistered.json').exists()
    assert all(not (EXP/d).exists() or not any((EXP/d).iterdir()) for d in ('protocol','graphs','apl-bundle','cases','gpu','oracle'))
    assert shutil.disk_usage(EXP.parent).free>30*1024**3
    for d in ('protocol','apl-bundle','graphs','cases','gpu','oracle','component-audit','regression','calibration','audit','extension-build','logs'):(EXP/d).mkdir(parents=True,exist_ok=True)
    can=ROOT/'datasets/male-cns-v1.0/processed/neuron-only';legacy=ROOT/'models/male-cns-lif';lp=legacy/'olfactory-grounding-v1/phase4l.yaml'
    lock=ROOT/'experiments/phase4j-mb-physiology-repair/calibration/locked-parameters.json'
    files=[AF_MODEL/'manifest.json',AF_EXP/'manifest.json',AF_EXP/'oracle/no_input/coefficients.json',
        legacy/'deterministic-runtime-v2/phase4p_revalidate.py',lp,lock,can/'neuron-index/unified-metadata.parquet',can/'gpu/csr-weights.int64.npy',
        legacy/'author-runtime-lab-v1/author_ext.cpp',legacy/'author-runtime-lab-v1/author_ext.cu',legacy/'author-runtime-lab-v1/gpu_backend.py',
        legacy/'author-runtime-lab-v1/oracle_stress.py',AD_EXP/'external-resources/model.py',AF_MODEL/'oracle_full.py']
    lroot=ROOT/'experiments/phase4l-olfactory-grounding'
    files += [lroot/'odor-encoder'/x for x in ('frozen-odor-panel.parquet','mapped-glomerular-vectors.npz','selection.json')]
    files += [lroot/'receptor-mapping/olfactory-mapping.parquet']
    prior=ROOT/'experiments/phase4ag-author-apl'
    files += [prior/'manifest.json',ROOT/'models/male-cns-lif/author-apl-lab-v1/manifest.json']
    save(EXP/'protocol/preregistered.json',dict(timestamp_utc=datetime.now(timezone.utc).isoformat(),sources=sources(),config=CFG,
        inputs=[dict(path=str(p.relative_to(ROOT)),sha256=digest(p)) for p in files],before_any_trajectory=True,
        parent_commit=subprocess.check_output(['git','-C',str(ROOT),'rev-parse','HEAD'],text=True).strip()))
    save(EXP/'regression/old-integrity-before.json',integrity())
    science_entry=ROOT/'models/male-cns-lif/.venv/bin/python'
    command=[str(science_entry),'-B','-c',
        'import sys,json,importlib.util; print(json.dumps(dict(executable=sys.executable,prefix=sys.prefix,base_prefix=sys.base_prefix,numpy_visible=importlib.util.find_spec("numpy") is not None,torch_visible=importlib.util.find_spec("torch") is not None)))']
    environment=json.loads(subprocess.check_output(command,text=True))
    assert environment['executable']==str(science_entry) and environment['prefix']==str(science_entry.parent.parent)
    assert environment['prefix']!=environment['base_prefix'] and environment['numpy_visible'] and environment['torch_visible']
    save(EXP/'protocol/execution-preflight.json',dict(command=command,observed=environment,
        before_any_trajectory=True,torch_imported=False,CUDA_initialized=False))
    locked=json.loads(lock.read_text());assert locked['parameters_locked'] and locked['calibration_only']
    assert {k:locked['selected'][k] for k in ('gain','normalization','tau_ms')}==CFG['apl']
    meta=pd.read_parquet(can/'neuron-index/unified-metadata.parquet').sort_values('neuron_idx').reset_index(drop=True)
    assert len(meta)==CFG['neurons'] and np.array_equal(meta.neuron_idx,np.arange(CFG['neurons']))
    cls=meta.annotation_class.fillna('').astype(str).to_numpy();typ=meta.annotation_type.fillna('').astype(str).to_numpy()
    kc=np.flatnonzero(cls=='Kenyon_Cell').astype(np.int32);pn=np.flatnonzero(cls=='ALPN').astype(np.int32);apl=np.flatnonzero(typ=='APL').astype(np.int32)
    assert len(kc)==4064 and len(apl)==2
    fg=AF_EXP/'graph';ptr=np.load(fg/'ptr.npy',mmap_mode='r');dst=np.load(fg/'dst.npy',mmap_mode='r');src=np.load(fg/'src.npy',mmap_mode='r')
    signed=np.load(fg/'signed_raw.npy',mmap_mode='r');base=np.load(AF_EXP/'oracle/no_input/weights.npy',mmap_mode='r');raw=np.load(can/'gpu/csr-weights.int64.npy',mmap_mode='r')
    assert len(dst)==CFG['edges'] and int(raw.sum())==CFG['raw_weight_sum']
    aplout=np.isin(src,apl);aplids=np.flatnonzero(aplout).astype(np.int64);assert len(aplids)==5845
    I=np.zeros((2,len(meta)),np.float32);O=np.zeros((2,len(kc)),np.float32)
    alook={int(v):i for i,v in enumerate(apl)};kl=np.full(len(meta),-1,np.int32);kl[kc]=np.arange(len(kc),dtype=np.int32)
    allowed=np.isin(src,np.r_[kc,pn])&np.isin(dst,apl)
    for eid in np.flatnonzero(allowed):I[alook[int(dst[eid])],int(src[eid])]+=float(raw[eid])
    for eid in np.flatnonzero(aplout&np.isin(dst,kc)):O[alook[int(src[eid])],kl[dst[eid]]]+=float(raw[eid])
    iscale=float(I[I>0].sum()/np.count_nonzero(np.any(I>0,axis=0)));oscale=float(O[O>0].mean());bundle=EXP/'apl-bundle'
    for name,value in dict(input_matrix=I/iscale,output_matrix=O/oscale,input_raw=I,output_raw=O,kc_globals=kc,pn_globals=pn,apl_globals=apl,apl_out_ids=aplids).items():arr(bundle/(name+'.npy'),value)
    save(bundle/'scales.json',dict(input_scale=iscale,output_scale=oscale))
    save(bundle/'provenance.json',dict(canonical_root=str(can),source=str(legacy/'deterministic-runtime-v2/phase4p_revalidate.py'),source_sha256=digest(legacy/'deterministic-runtime-v2/phase4p_revalidate.py'),
        parameters=CFG['apl'],lock_file=str(lock),lock_sha256=digest(lock),full_dimensions=[list(I.shape),list(O.shape)],
        APL=len(apl),KC=len(kc),ALPN=len(pn),APL_fast_edges=len(aplids),APL_nonKC_edges=int(np.count_nonzero(aplout&~np.isin(dst,kc))),
        APL_nonKC_contacts=int(raw[aplout&~np.isin(dst,kc)].sum()),normalization_recomputed_in_original_order=True,
        raw_files_modified=False))
    full=(ptr,dst,src,signed,base);fullnodes=graph(EXP/'graphs/full',np.arange(len(meta),dtype=np.int32),full)
    pscore=I.sum(axis=0);pnrank=pn[np.lexsort((pn,-pscore[pn]))];kscore=O.sum(axis=0);kcrank=kc[np.lexsort((kc,-kscore))]
    candidates=np.flatnonzero(np.isin(src,pn)&np.isin(dst,kc));best=int(candidates[np.argmax(raw[candidates])]);pairpn=int(src[best])
    pe=np.arange(ptr[pairpn],ptr[pairpn+1],dtype=np.int64);pe=pe[np.isin(dst[pe],kc)];pe=pe[np.lexsort((dst[pe],-raw[pe]))]
    collision_kc=np.unique(dst[pe[:3]]);assert len(collision_kc)==3
    nonkc=aplids[~np.isin(dst[aplids],kc)];nonkc_target=int(dst[nonkc[np.argmax(raw[nonkc])]])
    nodes=np.unique(np.r_[apl,pnrank[:16],kcrank[:16],pairpn,collision_kc,nonkc_target,[1197,1381,1511,82312,133666]]).astype(np.int32)
    fixturegraph=EXP/'graphs/fixtures';graph(fixturegraph,nodes,full)
    p=int(pnrank[0]);k=int(kcrank[0]);kc1,kc2,kc3=map(int,collision_kc)
    designs=[('zero',100,[],None),('pn_pulse',300,[(0,p)],None),
        ('kc_ref_boundary',100,[(t,k) for t in (0,1,21,22,23,44)],None),
        ('two_sources',300,[(t,n) for t in (0,22,44,66,88) for n in (p,k)],None),
        ('arrival_collision',120,[(0,pairpn),(0,kc1),(18,kc2),(21,kc1)],None),
        ('natural_request',100,[(0,k)],dict(v={k:-44.0})),
        ('apl_only',100,[(0,int(a)) for a in apl],None),
        ('high_fanin',400,[(t,int(n)) for t in (0,22,44) for n in nodes if pscore[n]>0],dict(g={k:.123}))]
    assert [v[0] for v in designs]==CFG['fixtures']
    for name,steps,requests,initial in designs:case(name,'fixture',steps,nodes,fixturegraph,requests,initial)
    meta.iloc[nodes][['neuron_idx','body_id','annotation_class','annotation_type','neurotransmitter']].to_csv(EXP/'protocol/fixture-identities.csv',index=False)
    save(EXP/'protocol/fixture-selection.json',dict(pn=p,kc=k,pair_PN=pairpn,arrival_KCs=collision_kc.tolist(),real_PN_KC_edge=best,
        nonKC_APL_target=nonkc_target,method='Static raw-count ranks and exact metadata; retain ALL existing induced edges and full APL matrices/norms',runtime_selection=False))
    os.environ['CUDA_VISIBLE_DEVICES']=CFG['gpu_uuid'];sys.path.insert(0,str(legacy));sys.path.insert(0,str(legacy/'deterministic-runtime-v2'))
    sys.path.insert(0,str(legacy/'terminal-scope-diagnostic-v1'));from diagnostic import queues
    queues()
    import torch
    from phase4p_revalidate import door_inputs,make_drive
    oc,panel,vectors,selection,mapping=door_inputs();cal=panel[panel.split=='calibration'].copy();assert len(panel)==12 and len(cal)==6
    assert oc['odor_encoder']['evaluation_seeds']==CFG['calibration']['seeds']
    assert float(oc['odor_encoder']['duration_ms'])==CFG['calibration']['duration_ms']
    cal.to_parquet(EXP/'protocol/calibration-panel.parquet',index=False);calrecords=[]
    for mode in CFG['calibration']['modes']:
        for oi,(_,odor) in enumerate(cal.iterrows()):
            targets,rates=make_drive(oc,panel,vectors,selection,mapping,odor.odor_id,mode)
            assert len(np.unique(targets))==len(targets),'Original encoder emits duplicate PN targets; audit instead of silently merging'
            rates=np.asarray(rates,np.float32);steps=round(CFG['calibration']['duration_ms']/.1)
            for seed in CFG['calibration']['seeds']:
                gen=torch.Generator(device='cuda');gen.manual_seed(seed)
                original_probability=rates*.1/1000.0
                explicit_probability=rates*np.float32(.1)/np.float32(1000.)
                assert original_probability.dtype==np.float32 and np.array_equal(original_probability.view(np.uint32),explicit_probability.view(np.uint32))
                probs=torch.as_tensor(original_probability,dtype=torch.float32,device='cuda')
                forced=(torch.rand((steps,len(targets)),generator=gen,device='cuda')<probs[None,:]).cpu().numpy()
                tt,cc=np.nonzero(forced);req=np.column_stack([tt,targets[cc]]).astype(np.int32)
                name=f'cal-{mode}-{oi:02d}-s{seed}'
                case(name,'calibration',steps,fullnodes,EXP/'graphs/full',req,extra=dict(odor_id=str(odor.odor_id),odor_name=str(odor.odor_name),
                    mode=mode,seed=seed,odor_index=oi,expected_external_events=float(rates.sum())*steps*.1/1000.,target_count=len(targets),rate_sum_hz=float(rates.sum())))
                with (EXP/'cases'/name/'pn-drive.npz').open('xb') as f:np.savez_compressed(f,targets=targets,rates_hz=rates)
                calrecords.append(dict(name=name,mode=mode,odor_index=oi,odor_id=str(odor.odor_id),seed=seed,requests=len(req),rates_sum=float(rates.sum())))
    save(EXP/'protocol/calibration-cases.json',calrecords)
    save(EXP/'protocol/encoder-audit.json',dict(panel_count=12,cal_count=6,median_mapped_response_sum=float(np.median(panel.mapped_response_sum)),
        target_rng_column_order='Original pn_drive, no sorting before torch.rand',seeds=CFG['calibration']['seeds'],
        requested_before_any_dynamics=True,formal_AB_run=False,heldout_run=False,torch=torch.__version__,numpy=np.__version__,
        generation='Original rates*.1/1000 float32 GPU Bernoulli, independent per-case original seed'))
    revision=[]
    for folder in sorted((EXP/'cases').iterdir()):
        with np.load(folder/'requests.npz',allow_pickle=False) as new,np.load(prior/'cases'/folder.name/'requests.npz',allow_pickle=False) as old:
            assert new.files==old.files and all(new[key].dtype==old[key].dtype and np.array_equal(new[key],old[key]) for key in old.files)
        revision.append(dict(case=folder.name,sparse_request_arrays_exact=True))
    save(EXP/'protocol/revision-input-equivalence.json',dict(prior_experiment=str(prior),cases=revision,
        input_arrays_unchanged=True,reason='Execution-only venv invocation fix; no trajectory outcome used for input choice'))
    save(EXP/'protocol/prepared-hashes.json',[dict(path=str(p.relative_to(EXP)),bytes=p.stat().st_size,sha256=digest(p))
        for d in ('graphs','apl-bundle','cases','protocol') for p in sorted((EXP/d).rglob('*')) if p.is_file() and p.name!='prepared-hashes.json'])
    budget();print('PREPARED8 fixtures and60 fixed CAL inputs; no trajectories; frozen integrity PASS',flush=True)

if __name__=='__main__':main()