"""Read-only canonical graph adaptation, before any neural trajectory."""
from datetime import datetime, timezone
import json
import platform
import shutil
import subprocess
import numpy as np
import pandas as pd
import yaml
from context import ROOT, MODEL, EXP, AD, AE_MODEL, AE_EXP, CFG, save, digest, sources, integrity, budget_check

def main():
    assert not (EXP/'protocol').exists()
    assert shutil.disk_usage(EXP.parent).free > 25*1024**3
    for name in ('protocol','graph','cases','oracle','gpu','regression','audit','logs','extension-build'):
        (EXP/name).mkdir(parents=True, exist_ok=False)
    can = ROOT/'datasets/male-cns-v1.0/processed/neuron-only'
    raw_paths = [can/'gpu/csr-indptr.int64.npy', can/'gpu/csr-indices.int32.npy',
        can/'gpu/csr-weights.int64.npy', can/'neuron-index/unified-metadata.parquet',
        ROOT/'models/male-cns-lif/runtime-v1/phase4c.yaml', AD/'external-resources/model.py',
        AD/'manifest.json', AE_MODEL/'manifest.json', AE_EXP/'manifest.json',
        AE_MODEL/'gpu_backend.py',AE_MODEL/'author_ext.cpp',AE_MODEL/'author_ext.cu',AE_MODEL/'oracle_stress.py']
    save(EXP/'protocol/preregistered.json', dict(timestamp_utc=datetime.now(timezone.utc).isoformat(),
        config=CFG, sources=sources(), before_model_execution=True,
        inputs=[dict(path=str(p.relative_to(ROOT)),sha256=digest(p)) for p in raw_paths],
        parent_commit=subprocess.check_output(['git','-C',str(ROOT),'rev-parse','HEAD'],text=True).strip()))
    save(EXP/'regression/old-integrity-before.json',integrity())
    ptr = np.load(raw_paths[0], mmap_mode='r'); dst = np.load(raw_paths[1], mmap_mode='r')
    raw = np.load(raw_paths[2], mmap_mode='r'); meta = pd.read_parquet(raw_paths[3]).sort_values('neuron_idx').reset_index(drop=True)
    n,e = CFG['neurons'],CFG['edges']
    assert ptr.dtype == np.int64 and dst.dtype == np.int32 and raw.dtype == np.int64
    assert len(meta)==n and np.array_equal(meta.neuron_idx,np.arange(n))
    assert ptr.shape==(n+1,) and ptr[0]==0 and ptr[-1]==e and np.all(np.diff(ptr)>=0)
    assert dst.shape == raw.shape == (e,) and dst.min()>=0 and dst.max()<n
    assert raw.min()>0 and int(raw.sum())==CFG['raw_weight_sum']
    policy = yaml.safe_load(raw_paths[4].read_text())['sign_policy']
    signs=[]; nt_counts={}
    for nt in meta.neurotransmitter:
        s='' if pd.isna(nt) else str(nt).strip().lower()
        signs.append(policy['unknown_sign'] if not s else policy['unclear_sign'] if s=='unclear' else policy['neurotransmitter_sign'][s])
        nt_counts[s or 'NULL'] = nt_counts.get(s or 'NULL',0)+1
    signs=np.asarray(signs,np.int8); assert set(signs).issubset({-1,0,1})
    src=np.repeat(np.arange(n,dtype=np.int32),np.diff(ptr)); signed=raw*signs[src]
    edge_id=np.arange(e,dtype=np.int64);graph=EXP/'graph'
    for name,array in dict(nodes=np.arange(n,dtype=np.int32),edge_id=edge_id,src=src,dst=dst,
                            signed_raw=signed,ptr=ptr,signs=signs).items():
        with (graph/(name+'.npy')).open('xb') as f: np.save(f,array,allow_pickle=False)
    pd.DataFrame({'local_index':np.arange(n,dtype=np.int32)},index=pd.Index(np.arange(n),name='global_neuron_idx')).to_csv(graph/'completeness.csv')
    pd.DataFrame({'Presynaptic_Index':src,'Postsynaptic_Index':dst,
                  'Excitatory x Connectivity':signed,'edge_id':edge_id},copy=False).to_parquet(graph/'connectivity.parquet',index=False)
    idx=np.bincount(dst,minlength=n);out=np.diff(ptr)
    assert int(np.count_nonzero(src==dst))==101 and int(np.count_nonzero((idx==0)&(out==0)))==217
    save(graph/'audit.json',dict(neurons=n,edges=e,raw_weight_sum=int(raw.sum()),
        self_loops=int(np.count_nonzero(src==dst)),isolates=int(np.count_nonzero((idx==0)&(out==0))),
        signs={str(int(s)):int(np.count_nonzero(signs==s)) for s in np.unique(signs)},neurotransmitters=nt_counts,
        canonical_edge_id='original CSR slot; source sorted; no reordering, filtering, normalization or new edge',
        body_mapping_input=str(raw_paths[3]),original_weights_modified=False))
    model=json.loads((AD/'protocol/preregistered.json').read_text())['config']['model']
    for trial in CFG['trials']:
        case=EXP/'cases'/trial['name'];case.mkdir()
        save(case/'spec.json',dict(**trial,dt_ms=model['dt_ms'],duration_ms=trial['steps']*model['dt_ms'],model=model,
             graph_dir=str(graph),precision=CFG['precision'],all_requests_at_tick=0,apl_overlay=False,training=False,
             budget_root=str(EXP),output_budget_bytes=CFG['new_output_budget_bytes']))
        meta.iloc[trial['seed_nodes']][['neuron_idx','body_id','annotation_type','neurotransmitter']].to_csv(case/'seed-identities.csv',index=False)
    save(EXP/'protocol/environment.json',dict(python=platform.python_version(),platform=platform.platform(),
        numpy=np.__version__,free_disk_bytes=shutil.disk_usage(EXP).free))
    save(EXP/'protocol/graph-case-hashes.json',[dict(path=str(p.relative_to(EXP)),sha256=digest(p),bytes=p.stat().st_size)
        for root in (graph,EXP/'cases') for p in sorted(root.rglob('*')) if p.is_file()])
    budget_check();print('PREPARED full canonical graph and3 fixed trials; frozen integrity PASS',flush=True)

if __name__=='__main__': main()