"""CPU-only, read-only APL/PN implementation scope audit. No scientific modules imported."""
import argparse
import ast
import concurrent.futures as cf
from datetime import datetime, timezone
import gzip
import hashlib
import json
from pathlib import Path
import resource
import time
import urllib.request
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import yaml


def digest(path):
    h=hashlib.sha256()
    with path.open('rb') as f:
        for b in iter(lambda:f.read(8<<20),b''):h.update(b)
    return h.hexdigest()


def save(path,value):
    path.parent.mkdir(parents=True,exist_ok=True)
    with path.open('x',encoding='utf-8') as f:
        json.dump(value,f,ensure_ascii=False,indent=2,sort_keys=True,allow_nan=False)
        f.write('\n')


def immutable_check(root,snapshot,workers):
    summary=json.loads((snapshot/'summary.json').read_text())
    assert digest(snapshot/'files.jsonl.gz')==summary['files_jsonl_gz_sha256']
    with gzip.open(snapshot/'files.jsonl.gz','rt') as f:rows=[json.loads(x) for x in f]
    def check(r):
        p=root/r['path']
        if r['kind']=='symlink':
            ok=p.is_symlink() and str(p.readlink())==r['target']
        else:
            ok=p.is_file() and p.stat().st_size==r['bytes'] and digest(p)==r['sha256']
        return None if ok else r['path']
    with cf.ThreadPoolExecutor(max_workers=workers) as pool:bad=[x for x in pool.map(check,rows) if x]
    return {'pass':not bad,'checked':len(rows),'mismatches':bad,'reference_sha256':summary['files_jsonl_gz_sha256'],
            'exclusions':summary['excluded_directory_names'],'new_stage_additions_allowed':True,
            'meaning':'All original listed bytes checked; only newly named stage outputs are additions.'}


def population(meta):
    out=meta.annotation_class.fillna('UNKNOWN_CLASS').astype(str).to_numpy(object)
    for typ in ('APL','DPM'):out[meta.annotation_type.eq(typ).to_numpy()]=typ
    return out


def edge_frame(ids,sources,targets,weights,meta,pops):
    src=sources[ids];dst=targets[ids]
    f=pd.DataFrame({'edge_id':ids,'source_neuron_idx':src,'target_neuron_idx':dst,'raw_synapse_count':weights[ids]})
    for pref,arr in [('source',src),('target',dst)]:
        f[pref+'_population']=pops[arr]
        for col in ('body_id','annotation_class','annotation_type','annotation_instance','annotation_somaSide','neurotransmitter','nt_ground_truth','annotation_receptorType'):
            f[pref+'_'+col]=meta[col].to_numpy()[arr]
    return f


def raw_incident(path,bodyids):
    tables=[]
    with pa.memory_map(str(path),'r') as f:
        reader=pa.ipc.open_file(f)
        assert reader.schema.names==['body_pre','body_post','weight'],str(reader.schema)
        ids=pa.array(bodyids.astype(np.int64))
        for i in range(reader.num_record_batches):
            batch=reader.get_batch(i)
            keep=pc.or_(pc.is_in(batch.column(0),value_set=ids),pc.is_in(batch.column(1),value_set=ids))
            selected=batch.filter(keep)
            if len(selected):tables.append(pa.Table.from_batches([selected]))
    return pa.concat_tables(tables).to_pandas()


def official_head(timeout):
    base='https://storage.googleapis.com/flyem-male-cns/v1.0/connectome-data/flat-connectome/'
    out=[]
    for name in ('syn-partners-male-cns-v1.0-minconf-0.5.feather','syn-points-male-cns-v1.0-minconf-0.5.feather'):
        row={'url':base+name,'retrieved_utc':datetime.now(timezone.utc).isoformat(),'method':'HEAD','body_downloaded':False}
        try:
            with urllib.request.urlopen(urllib.request.Request(base+name,method='HEAD'),timeout=timeout) as r:
                row.update(status=r.status,content_length=int(r.headers['Content-Length']),etag=r.headers.get('ETag'),
                           generation=r.headers.get('x-goog-generation'),upstream_hash_header=r.headers.get('x-goog-hash'),
                           last_modified=r.headers.get('Last-Modified'),accept_ranges=r.headers.get('Accept-Ranges'))
        except Exception as e:row.update(status='UNAVAILABLE',error=type(e).__name__+': '+str(e))
        out.append(row)
    return out


def seal(root):
    paths=sorted(p for p in root.rglob('*') if p.is_file() and p.name not in ('manifest.json','checksums.sha256'))
    rows=[{'path':p.relative_to(root).as_posix(),'bytes':p.stat().st_size,'sha256':digest(p)} for p in paths]
    save(root/'manifest.json',{'files':rows,'created_utc':datetime.now(timezone.utc).isoformat()})
    with (root/'checksums.sha256').open('x') as f:
        for r in rows:f.write(r['sha256']+'  '+r['path']+'\n')
    return len(rows)


def run(cfg):
    pa.set_cpu_count(4);pa.set_io_thread_count(4)
    start=time.perf_counter();root=Path(cfg['root']);exp=Path(cfg['experiment_root']);model=Path(cfg['model_root'])
    assert exp.is_relative_to(root/'experiments') and model.is_relative_to(root/'models/male-cns-lif')
    if exp.exists() and any(exp.iterdir()):raise RuntimeError('Refusing nonempty experiment directory')
    exp.mkdir(parents=True,exist_ok=True)
    for d in ('protocol','audit','connectivity','metadata','spatial','external-resources','regression','performance'):(exp/d).mkdir()
    save(exp/'protocol/preregistered.json',{'config':cfg,'source_hashes':{p.name:digest(p) for p in model.iterdir() if p.is_file()},
        'created_utc':datetime.now(timezone.utc).isoformat(),'before_analysis':True})
    pre=immutable_check(root,root/cfg['snapshot'],cfg['scope']['snapshot_workers']);save(exp/'regression/preflight.json',pre)
    assert pre['pass'],'original bytes differ'
    canonical=root/cfg['canonical'];prior=root/cfg['prior'];expect=cfg['expected']
    meta=pd.read_parquet(canonical/'neuron-index/unified-metadata.parquet').sort_values('neuron_idx').reset_index(drop=True)
    ptr=np.load(canonical/'gpu/csr-indptr.int64.npy',mmap_mode='r');dst=np.load(canonical/'gpu/csr-indices.int32.npy',mmap_mode='r');w=np.load(canonical/'gpu/csr-weights.int64.npy',mmap_mode='r')
    assert len(meta)==expect['neurons'] and np.array_equal(meta.neuron_idx,np.arange(len(meta)))
    assert str(meta.neuron_idx.dtype)=='int32' and str(meta.body_id.dtype)=='uint64' and meta.body_id.is_unique
    assert len(ptr)==len(meta)+1 and ptr[0]==0 and ptr[-1]==len(w)==len(dst)==expect['edges']
    assert ptr.dtype==np.int64 and dst.dtype==np.int32 and w.dtype==np.int64
    assert int(w.sum())==expect['synapse_count_sum'] and np.all(np.diff(ptr)>=0) and w.min()>0
    src=np.repeat(np.arange(len(meta),dtype=np.int32),np.diff(ptr));pops=population(meta)
    apl=np.flatnonzero(meta.annotation_type.eq('APL'));kc=np.flatnonzero(meta.annotation_class.eq('Kenyon_Cell'));pn=np.flatnonzero(meta.annotation_class.eq('ALPN'))
    assert len(apl)==expect['APL'] and len(kc)==expect['KC']
    oi=np.flatnonzero(np.isin(src,apl));ii=np.flatnonzero(np.isin(dst,apl))
    outgoing=edge_frame(oi,src,dst,w,meta,pops);incoming=edge_frame(ii,src,dst,w,meta,pops)
    outgoing['fast_current_masked']=True;outgoing['graded_output_covered']=np.isin(dst[oi],kc)
    incoming['graded_input_used']=np.isin(src[ii],np.r_[kc,pn])
    # This describes the graded-state input only. Fast synapses into APL remain in CSR/runtime.
    watched=json.loads((prior/'protocol/watched.json').read_text())['all_incoming_ALPNs']
    focus=outgoing[outgoing.target_neuron_idx.isin(watched)].copy()
    old=pd.read_parquet(prior/'audit/APL-target-coverage.parquet')
    fields=['edge_id','source_neuron_idx','target_neuron_idx','raw_synapse_count']
    assert np.array_equal(outgoing[fields].sort_values('edge_id').to_numpy(np.int64),old[fields].sort_values('edge_id').to_numpy(np.int64))
    nonkc=outgoing[~outgoing.graded_output_covered]
    assert (len(outgoing),int(outgoing.raw_synapse_count.sum()),len(nonkc),int(nonkc.raw_synapse_count.sum()),len(focus),int(focus.raw_synapse_count.sum()))==(expect['APL_outgoing_edges'],expect['APL_outgoing_contacts'],expect['non_KC_edges'],expect['non_KC_contacts'],expect['watched_PN_edges'],expect['watched_PN_contacts'])
    print('CSR and Phase4Z coverage exact',flush=True)
    raw_path=root/'datasets/male-cns-v1.0/raw/connectivity/connectome-weights-male-cns-v1.0-minconf-0.5.feather'
    raw=raw_incident(raw_path,meta.iloc[apl].body_id.to_numpy())
    valid=raw.body_pre.isin(meta.body_id)&raw.body_post.isin(meta.body_id)
    actual=raw.loc[valid].sort_values(['body_pre','body_post']).reset_index(drop=True)
    assert not actual.duplicated(['body_pre','body_post']).any()
    incident=pd.concat([outgoing,incoming]).drop_duplicates('edge_id')
    expected=incident[['source_body_id','target_body_id','raw_synapse_count']].rename(columns={'source_body_id':'body_pre','target_body_id':'body_post','raw_synapse_count':'weight'}).sort_values(['body_pre','body_post']).reset_index(drop=True)
    assert np.array_equal(actual.to_numpy(np.int64),expected.to_numpy(np.int64))
    raw.to_parquet(exp/'connectivity/raw-APL-incident-including-noncanonical.parquet',index=False)
    print('Raw Feather APL incident edges exact',len(actual),flush=True)
    outgoing.to_parquet(exp/'connectivity/APL-outgoing.parquet',index=False);incoming.to_parquet(exp/'connectivity/APL-incoming.parquet',index=False)
    for label,frame,col in [('output',outgoing,'target_population'),('input',incoming,'source_population')]:
        grouped=frame.groupby(col,dropna=False).agg(edges=('edge_id','size'),contacts=('raw_synapse_count','sum')).reset_index()
        grouped.to_parquet(exp/'connectivity'/f'APL-{label}-population-summary.parquet',index=False)
        grouped.to_csv(exp/'connectivity'/f'APL-{label}-population-summary.csv',index=False)
    ignored=incoming[~incoming.graded_input_used]
    ignored.to_parquet(exp/'connectivity/graded-input-not-used.parquet',index=False)
    pairs=incoming[['source_neuron_idx','target_neuron_idx','raw_synapse_count','edge_id']].rename(columns={'source_neuron_idx':'target_neuron_idx','target_neuron_idx':'source_neuron_idx','raw_synapse_count':'reverse_contacts','edge_id':'reverse_edge_id'})
    pn_out=outgoing[outgoing.target_neuron_idx.isin(pn)].merge(pairs,on=['source_neuron_idx','target_neuron_idx'],how='left',validate='one_to_one')
    pn_out['has_PN_to_APL']=pn_out.reverse_edge_id.notna()
    pn_out['watched_in_Phase4Z']=pn_out.target_neuron_idx.isin(watched)
    pn_out.to_parquet(exp/'connectivity/APL-PN-reciprocity.parquet',index=False)
    pn_out.to_csv(exp/'connectivity/APL-PN-reciprocity.csv',index=False)
    downstream=[]
    for node in pn_out.target_neuron_idx.unique():
        ids=np.arange(ptr[node],ptr[node+1]);cl=pops[dst[ids]]
        for pop in ('Kenyon_Cell','ALLN','APL'):
            selected=ids[cl==pop]
            downstream.append({'PN_neuron_idx':int(node),'PN_body_id':int(meta.iloc[node].body_id),'PN_type':meta.iloc[node].annotation_type,
                               'target_population':pop,'edges':len(selected),'contacts':int(w[selected].sum())})
    pd.DataFrame(downstream).to_parquet(exp/'connectivity/APL-target-PN-downstream.parquet',index=False)
    focus=pn_out[pn_out.watched_in_Phase4Z].copy()
    focus[['source_neuron_idx','source_body_id','source_annotation_instance','target_neuron_idx','target_body_id','target_annotation_type','raw_synapse_count','reverse_contacts','reverse_edge_id']].to_csv(exp/'connectivity/watched-PN-reciprocity.csv',index=False)
    # Reconstruct the exact frozen matrices on CPU; do not simulate or apply any weights.
    inp=np.zeros((len(apl),len(meta)),np.float32);output=np.zeros((len(apl),len(kc)),np.float32)
    kmap={int(n):i for i,n in enumerate(kc)};amap={int(n):i for i,n in enumerate(apl)}
    for r in incoming[incoming.graded_input_used].itertuples():inp[amap[r.target_neuron_idx],r.source_neuron_idx]+=r.raw_synapse_count
    for r in outgoing[outgoing.graded_output_covered].itertuples():output[amap[r.source_neuron_idx],kmap[r.target_neuron_idx]]+=r.raw_synapse_count
    input_scale=float(inp[inp>0].sum()/np.count_nonzero(np.any(inp>0,axis=0)));output_scale=float(output[output>0].mean())
    coverage={'input_scale':input_scale,'output_scale':output_scale,'input_matrix_shape':list(inp.shape),'output_matrix_shape':list(output.shape),
              'KC_targets_with_graded_input':int(np.count_nonzero(np.any(output>0,axis=0))),
              'normalized_input_matrix_sha256':hashlib.sha256((inp/input_scale).tobytes()).hexdigest(),
              'normalized_output_matrix_sha256':hashlib.sha256((output/output_scale).tobytes()).hexdigest()}
    save(exp/'audit/frozen-matrix-reconstruction.json',coverage)
    frozen=json.loads((prior/'audit/APL-target-coverage.json').read_text());source=Path(frozen['source'])
    assert digest(source)==frozen['source_sha256']
    body=source.read_text();tree=ast.parse(body)
    cls=next(n for n in tree.body if isinstance(n,ast.ClassDef) and n.name=='DeterministicAPLSim')
    with (exp/'audit/frozen-APL-class-source.txt').open('x') as f:f.write(ast.get_source_segment(body,cls)+'\n')
    save(exp/'audit/frozen-implementation-provenance.json',{'path':str(source),'sha256':digest(source),'same_as_Phase4Z':True,'models_imported':False})
    ids=np.unique(np.r_[apl,pn_out.target_neuron_idx]);meta.iloc[ids].to_parquet(exp/'metadata/APL-and-target-PNs.parquet',index=False)
    spatial=pn_out[['edge_id','source_neuron_idx','source_body_id','target_neuron_idx','target_body_id','target_annotation_type','raw_synapse_count']].copy()
    for col in ('synapse_location','synapse_ROI','PN_subcellular_target','postsynaptic_GABA_receptor','effective_inhibition_coefficient'):spatial[col]='unknown'
    spatial['reason']='Local body-pair counts and SWC centerlines do not contain synapse partner locations or target receptor measurements'
    spatial.to_parquet(exp/'spatial/APL-PN-unknowns.parquet',index=False)
    swc=[]
    for node in np.r_[apl,watched]:
        p=Path(meta.iloc[node].swc_path)
        a=np.loadtxt(p,comments='#',ndmin=2)
        assert a.shape[1]==7 and np.all(np.isfinite(a))
        swc.append({'neuron_idx':int(node),'body_id':int(meta.iloc[node].body_id),'path':str(p),'sha256':digest(p),
                    'nodes':len(a),'columns':int(a.shape[1]),'synaptic_partner_identity_in_standard_SWC':False,
                    'coordinates_used_to_infer_synapse_locations':False})
    save(exp/'spatial/watched-SWC-audit.json',swc)
    schemas=[]
    for p in sorted((root/'datasets/male-cns-v1.0/raw').glob('*/*.feather')):
        with pa.memory_map(str(p),'r') as f:sc=pa.ipc.open_file(f).schema
        schemas.append({'path':str(p),'sha256':digest(p),'fields':[{'name':x.name,'dtype':str(x.type)} for x in sc]})
    save(exp/'spatial/local-schemas.json',schemas)
    save(exp/'external-resources/official-data-access.json',official_head(cfg['scope']['official_network_timeout_s']))
    evidence={
      'official_definition':{'url':'https://male-cns.janelia.org/download/','dataset':'male-cns:v1.0','license':'CC-BY-4.0',
        'license_url':'https://creativecommons.org/licenses/by/4.0/','observation':'Official page specifies paired synapse coordinates/primary_post in syn-partners; local flat weights contain none.',
        'required_subset':'All outgoing APL partner records; exact body IDs from metadata; minimum confidence 0.5; include pre/post x/y/z, body IDs and primary_post.',
        'no_EM_needed':True,'coordinate_units':'synapse coordinates and original SWC in 8nm voxels; precomputed skeletons have different units'},
      'primary_physiology':[
        {'doi':'10.7554/eLife.74172','url':'https://elifesciences.org/articles/74172','claim':'Local APL inhibition at PN boutons and KC dendritic claws was studied; no edge-specific MaleCNS coefficient follows.'},
        {'doi':'10.7554/eLife.56954','url':'https://pmc.ncbi.nlm.nih.gov/articles/PMC7541083/','claim':'APL activity/inhibition can be spatially localized; global point-neuron inhibition is a modeling assumption.'}],
      'google_research':{'url':'https://www.research.google/blog/a-connectomics-milestone-mapping-the-complete-male-fruit-fly-brain/',
          'role':'Primary project provenance, not evidence of associative-learning performance'},
      'access_choices':['Authorized neuPrint male-cns:v1.0 targeted query, requiring user-configured API token',
                        'Explicitly approved official syn-partners bulk download into a NEW external-reference directory'],
      'not_attempted':['bulk synapse download','EM download','credential discovery','unauthenticated query bypass','new dynamics','GPU regression']}
    save(exp/'external-resources/evidence-and-limitations.json',evidence)
    results={'phase':'4AA','status':'PASS-STATIC-SCOPE-AUDIT / BLOCKED-SPATIAL-PHYSIOLOGY','candidate':False,
       'neurons':len(meta),'graph_edges':len(w),'synapse_count_sum':int(w.sum()),'APL_neurons':len(apl),'KC_neurons':len(kc),'ALPN_neurons':len(pn),
       'APL_outgoing_edges':len(outgoing),'APL_outgoing_contacts':int(outgoing.raw_synapse_count.sum()),
       'non_KC_uncovered_edges':len(nonkc),'non_KC_uncovered_contacts':int(nonkc.raw_synapse_count.sum()),'non_KC_uncovered_targets':int(nonkc.target_neuron_idx.nunique()),
       'APL_to_PN_edges':len(pn_out),'APL_to_PN_contacts':int(pn_out.raw_synapse_count.sum()),'APL_target_PNs':int(pn_out.target_neuron_idx.nunique()),
       'APL_PN_pairs_with_reverse_edge':int(pn_out.has_PN_to_APL.sum()),'watched_PN_contacts':int(focus.raw_synapse_count.sum()),
       'APL_incoming_edges':len(incoming),'APL_incoming_contacts':int(incoming.raw_synapse_count.sum()),
       'incoming_edges_not_used_by_graded_state':len(ignored),'incoming_contacts_not_used_by_graded_state':int(ignored.raw_synapse_count.sum()),
       'raw_APL_incident_canonical_edges_checked':len(actual),'raw_APL_incident_noncanonical_rows':int((~valid).sum()),
       'raw_to_CSR_exact':True,'Phase4Z_static_coverage_exact':True,'APL_PN_pairs_with_confirmed_local_synapse_ROI':0,
       'APL_PN_pairs_with_confirmed_subcellular_target':0,'whole_cell_inhibition_repair_authorized_by_evidence':False,
       'new_simulation':False,'new_scientific_regression_run':False,'old_reference_files_rehashed':True,
       'weights_changed':False,'gpu_used':False,'ComfyUI_models_unloaded':False,'top_level_gates_passed':5,'top_level_gates_total':10}
    post=immutable_check(root,root/cfg['snapshot'],cfg['scope']['snapshot_workers']);save(exp/'regression/postflight.json',post);assert post['pass']
    results['frozen_integrity_pass']=post['pass'];save(exp/'audit/results.json',results)
    save(exp/'performance/resources.json',{'wall_seconds':time.perf_counter()-start,'max_RSS_KiB':resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,'gpu_used':False,'includes_two_full_integrity_scans':True})
    print(json.dumps(results,indent=2),flush=True)


if __name__=='__main__':
    ap=argparse.ArgumentParser();ap.add_argument('--config',type=Path,required=True);args=ap.parse_args()
    run(yaml.safe_load(args.config.read_text()))