"""Read-only frozen-brain timing audit. All outputs are confined to Phase 4R."""
from pathlib import Path
import argparse
import hashlib
import json
import os
import resource
import sys
import time

import numpy as np
import pandas as pd
import torch
import yaml

import phase4p
import phase4p_finalize
from phase4p_revalidate import DeterministicAPLSim, door_inputs, make_drive
from phase4h_audit import verify_checksum
from phase4q import last_spike_trace, edge_elig, random_us_steps


def save(p, obj):
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(obj, indent=2, sort_keys=True, allow_nan=False) + '\n')


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


def arrhash(*a):
    h=hashlib.sha256()
    for x in a: h.update(np.ascontiguousarray(x).tobytes())
    return h.hexdigest()


def quick_integrity(cfg):
    specs=[('P',Path(cfg['phase4p_root']),'manifests/checksums.sha256'),
           ('P_candidate',Path(cfg['runtime_root'])/'candidate','checksums.sha256'),
           ('Q',Path(cfg['phase4q_root']),'manifests/checksums.sha256'),
           ('Q_model',Path(cfg['phase4q_model']),'checksums.sha256')]
    return {k:verify_checksum(root,root/rel) for k,root,rel in specs}


def oracle(steps,nodes,lookup,us,window,tau,dt):
    # Independent oracle: sort event times separately for every KC, binary-search at US.
    out=np.zeros(int(lookup.max())+1,np.float32)
    for node in np.unique(nodes):
        ts=np.sort(steps[nodes==node]); pos=np.searchsorted(ts,us,side='right')-1
        if pos>=0:
            age=np.float32(us-int(ts[pos]))*np.float32(dt)
            if age<=np.float32(window): out[lookup[node]]=np.exp(-age/np.float32(tau))
    return out


def curves(run,lookup,multiplicity,windows,dt,nsteps,shift=0):
    last=np.full((nsteps,len(multiplicity)),-1,np.int32)
    steps=run['kc_event_step']+shift; nodes=run['kc_event_node'];valid=steps<nsteps
    last[steps[valid],lookup[nodes[valid]]]=steps[valid]
    np.maximum.accumulate(last,axis=0,out=last)
    age=(np.arange(nsteps,dtype=np.int32)[:,None]-last).astype(np.float32)*np.float32(dt)
    rows=[]
    for w in windows:
        e=np.zeros(age.shape,np.float32);ok=(last>=0)&(age<=np.float32(w))
        e[ok]=np.exp(-age[ok]/np.float32(w/3))
        rows.append(np.sum(e*multiplicity[None,:],axis=1,dtype=np.float64))
    return rows


def run_audit(cfg,exp,model):
    for name in ('audit','protocol','events','timing','fixed-event','network-effects','regression','performance','figures','manifests'):
        (exp/name).mkdir(parents=True,exist_ok=True)
    if (exp/'audit/preflight.json').exists(): raise RuntimeError('Existing run: refusing to overwrite. Review it explicitly.')
    pre=quick_integrity(cfg);save(exp/'audit/preflight.json',pre)
    if not all(x['passed'] for x in pre.values()): raise RuntimeError('Frozen P/Q integrity failed')
    pcfg=yaml.safe_load((Path(cfg['runtime_root'])/'phase4p.yaml').read_text())
    qcfg=yaml.safe_load((Path(cfg['phase4q_model'])/'phase4q.yaml').read_text())
    ocfg,panel,vectors,selection,mapping=door_inputs()
    cal=panel[panel.split=='calibration'].iloc[:2]
    assert not set(cal.odor_id)&{selection['ODOR_A']['odor_id'],selection['ODOR_B']['odor_id']}
    inputs=[];drives={};nsteps=round(cfg['duration_ms']/pcfg['model']['dt_ms']);dt=pcfg['model']['dt_ms']
    for oi,(_,odor) in enumerate(cal.iterrows()):
        label='CAL_'+chr(65+oi); t,r=make_drive(ocfg,panel,vectors,selection,mapping,str(odor.odor_id),'physiological');drives[label]=(t,r)
        pd.DataFrame({'neuron_idx':t,'rate_hz':r}).to_parquet(exp/'protocol'/f'{label}-PN-drive.parquet',index=False)
        for split in ('calibration','heldout'):
            for seed in cfg[split+'_seeds']:
                actual=seed+oi*100
                qs=int(random_us_steps(cfg['random_q_seed_base'],actual,1,nsteps)[0])
                ns=round(float(np.random.default_rng(cfg['random_n_seed_base']+actual).uniform(30,110))/dt)
                inputs.append({'label':label,'odor_id':str(odor.odor_id),'odor_name':str(odor.odor_name),
                  'split':split,'seed':seed,'actual_seed':actual,'matched_step':round(cfg['matched_us_ms']/dt),
                  'q_random_step':qs,'n_random_step':ns,'target_count':len(t),'rate_sum_hz':float(r.sum()),'drive_hash':arrhash(t,r)})
    save(exp/'protocol/preregistered-inputs.json',{'config':cfg,'inputs':inputs,'formal_AB_used':False,
         'saved_before_simulation':True,'selected_timing':None})
    sourcefiles=[Path(cfg['runtime_root'])/n for n in ('phase4p.py','phase4p_ext.cpp','phase4p_ext.cu','phase4p.yaml','phase4p_revalidate.py')]
    sourcefiles += [Path(cfg['phase4q_model'])/n for n in ('phase4q.py','phase4q.yaml')]
    sourcefiles += list(model.glob('phase4r.*'))
    save(exp/'protocol/source-hashes.json',{str(p):digest(p) for p in sourcefiles})
    static_findings={
      'odor_forcing':'constant PN rate vectors over all 120 ms, Bernoulli events per timestep; no brief CS pulse in this runtime',
      'matched_US':'one forced spike per selected DAN at 90 ms; external teaching gate sampled at same emission timestep',
      'random_Q':'one discrete uniform step over [0,120) ms; unchanged selected DANs/amplitude',
      'random_N':'continuous uniform [30,110] ms then rounded to a timestep',
      'Q_KC_timing':'KC soma/emission event time is called terminal time; does not add 18-step anatomical delay',
      'equal_delay_note':'shifting BOTH KC and US by 18 steps preserves relative timing; shifting only KC changes the teaching contract',
      'Q_control_scope':'D0/unrelated/far-before/far-after are algebraic eligibility tests, not physiological control simulations',
      'Q_ten_repeats':'same NumPy function called ten times on a synthetic event, not full-brain runtime repeats',
      'Q_checkpoint':'reconstructs trace from in-memory last-spike array, no disk serialize/restore',
      'Q_timing_matrix':'synchronous hypothetical activity on all in-compartment terminals, not measured odor-driven edge-group credit',
      'Q_future_far_case':'for window 250 ms, synthetic far-before KC step is negative; this is not an actual runtime event',
      'Q_independence':'calibration split selected independently of formal odors; frozen-stimuli file was written after simulations',
      'Q_comparison':'matched and random timing were run separately, so KC streams may differ through DAN/network feedback',
      'Q_claim_limit':'calibration failure does not prove universal credit-rule failure or associative-memory failure'}
    save(exp/'audit/static-protocol-findings.json',static_findings)
    rt=phase4p.CanonicalRuntime(pcfg,Path(cfg['runtime_root']),exp);sim=DeterministicAPLSim(rt)
    assert rt.n==166700 and len(rt.indices_np)==25582938
    identity=sim.identity;comp=identity.compartment.astype(str).to_numpy()==cfg['target_compartment']
    lookup=np.full(sim.n,-1,np.int32);lookup[sim.kc]=np.arange(len(sim.kc),dtype=np.int32)
    source_local=lookup[identity.source_neuron_idx.to_numpy(np.int32)]
    assert np.all(source_local>=0)
    multi=np.bincount(source_local[comp],minlength=len(sim.kc)).astype(np.float32)
    dan=sim.dan_targets(cfg['matched_dan_type'],cfg['target_compartment'])
    identity.to_parquet(exp/'protocol/frozen-edge-identity.parquet',index=False)
    pd.DataFrame({'neuron_idx':dan,'body_id':sim.meta.body_id.to_numpy()[dan]}).to_parquet(exp/'protocol/DAN-targets.parquet',index=False)
    ids=identity.edge_id.to_numpy(np.int64)
    assert np.array_equal(sim.sources[ids],identity.source_neuron_idx.to_numpy())
    assert np.array_equal(rt.indices_np[ids],identity.target_neuron_idx.to_numpy())
    # Real DAN->KC/MBON edges document arrival provenance; no inferred compartment relations.
    dk=np.isin(sim.sources,dan)&np.isin(rt.indices_np,np.r_[sim.kc,sim.mbon])
    realids=np.flatnonzero(dk)
    pd.DataFrame({'edge_id':realids,'source_neuron_idx':sim.sources[realids],'target_neuron_idx':rt.indices_np[realids],
      'raw_synapse_count':rt.raw_np[realids],'delay_steps':rt.delay_steps}).to_parquet(exp/'protocol/real-DAN-to-KC-MBON-edges.parquet',index=False)
    telemetry=[];hist=[];fixed=[];grid=[];effects=[];edge_rows=[];oracle_rows=[];q_repro=[];repeat_rows=[]
    qold=pd.read_parquet(Path(cfg['phase4q_root'])/'calibration/frozen-stimuli.parquet')
    qmatrix=pd.read_parquet(Path(cfg['phase4q_root'])/'calibration/window-matrix.parquet')
    walls=[];t0=time.perf_counter();cpu0=resource.getrusage(resource.RUSAGE_SELF);torch.cuda.reset_peak_memory_stats()
    for ix,spec in enumerate(inputs):
        tag=f"{spec['split']}-{spec['label']}-{spec['seed']}";targets,rates=drives[spec['label']]
        runs={}; gates={'no_external_US':None,'matched_90ms':spec['matched_step'],
                        'random_Q_0_120ms':spec['q_random_step'],'random_N_30_110ms':spec['n_random_step']}
        for condition in cfg['conditions']:
            u=gates[condition];run=sim.run(targets,rates,spec['actual_seed'],cfg['duration_ms'],
              dan_targets=dan,dan_pulses=[] if u is None else [u*dt],collect_events=True)
            assert run['finite'];runs[condition]=run;walls.append(run['wall_seconds'])
            ev=run['kc_event_step'];nodes=run['kc_event_node'];loc=lookup[nodes];incomp=multi[loc]>0
            frame=pd.DataFrame({'emission_step':ev,'arrival_step':ev+rt.delay_steps,'neuron_idx':nodes,
                               'body_id':sim.meta.body_id.to_numpy()[nodes]})
            frame.to_parquet(exp/'events'/f'{tag}-{condition}.parquet',index=False)
            np.savez_compressed(exp/'events'/f'{tag}-{condition}-state.npz',counts=run['counts'],v=run['final_v'],g=run['final_g'])
            telemetry.append({**spec,'condition':condition,'external_US_step':u,'external_US_event_count':0 if u is None else len(dan),
              'KC_events':len(ev),'KC_active':int(np.count_nonzero(run['kc_counts'])),'compartment_source_events':int(incomp.sum()),
              'KC_fraction_first_30ms':float(np.mean(ev<300)) if len(ev) else 0.,
              'compartment_fraction_first_30ms':float(np.mean(ev[incomp]<300)) if incomp.any() else 0.,
              'compartment_events_85_to_90ms':int(np.count_nonzero(incomp&(ev>=850)&(ev<=900))),
              'KC_events_sha256':arrhash(ev,nodes),'wall_seconds':run['wall_seconds']})
            for begin in range(0,nsteps,50):
                select=(ev>=begin)&(ev<begin+50)
                hist.append({**spec,'condition':condition,'start_ms':begin*dt,'KC_spikes':int(select.sum()),
                    'compartment_source_spikes':int((select&incomp).sum()),
                    'plastic_terminal_emissions':int(multi[loc[select]].sum()),
                    'plastic_terminal_arrivals':int(multi[loc[(ev+rt.delay_steps>=begin)&(ev+rt.delay_steps<begin+50)]].sum())})
            if u is not None:
              for w in cfg['windows_ms']:
                tr,last=last_spike_trace(ev,nodes,lookup,u,w,w/3,dt)
                independent=oracle(ev,nodes,lookup,u,w,w/3,dt)
                oracle_rows.append({'tag':tag,'condition':condition,'window_ms':w,'exact':bool(np.array_equal(tr,independent))})
                assert np.array_equal(tr,independent)
                if spec['split']=='calibration' and condition in ('matched_90ms','random_Q_0_120ms'):
                    qc='matched' if condition=='matched_90ms' else 'random'
                    old=qmatrix[(qmatrix.label==spec['label'])&(qmatrix.seed==spec['seed'])&(qmatrix.condition==qc)&(qmatrix.window_ms==w)].iloc[0]
                    oldh=qold[(qold.label==spec['label'])&(qold.seed==spec['seed'])].iloc[0][qc+'_event_hash']
                    actual=float(edge_elig(tr,source_local,comp).sum())
                    q_repro.append({'tag':tag,'condition':condition,'window_ms':w,'event_hash_exact':oldh==arrhash(ev,nodes),
                                    'credit_exact':actual==float(old.total_eligibility),'old_credit':float(old.total_eligibility),'new_credit':actual})
        baseline=runs['no_external_US'];bev=baseline['kc_event_step'];bn=baseline['kc_event_node']
        # Dense arrays here are CPU accounting only; runtime spike outputs remain events.
        soma=curves(baseline,lookup,multi,cfg['windows_ms'],dt,nsteps)
        arrival=curves(baseline,lookup,multi,cfg['windows_ms'],dt,nsteps,rt.delay_steps)
        for wi,w in enumerate(cfg['windows_ms']):
            c=soma[wi];ca=arrival[wi]
            expected_q=float(c.mean());expected_n=float((c[300]/2+c[301:1100].sum()+c[1100]/2)/800)
            row={**spec,'window_ms':w,'matched_same_events':float(c[900]),'random_q_exact_expectation':expected_q,
                'random_n_exact_expectation':expected_n,'ratio_q':float(c[900]/expected_q) if expected_q else None,
                'ratio_n':float(c[900]/expected_n) if expected_n else None,'peak_ms_diagnostic_only':float(np.argmax(c)*dt),
                'matched_KC_arrival_US_emission':float(ca[900]),
                'matched_both_arrival':float(ca[900+rt.delay_steps]),
                'symmetric_delay_exact':bool(c[900]==ca[900+rt.delay_steps])}
            fixed.append(row)
            for step in range(nsteps):
                grid.append({'label':spec['label'],'seed':spec['seed'],'split':spec['split'],'window_ms':w,
                  'US_step':step,'credit_same_KC_events':float(c[step]),'KC_arrival_credit_US_emission':float(ca[step])})
            for condition in cfg['conditions'][1:]:
                u=gates[condition];run=runs[condition]
                own,_=last_spike_trace(run['kc_event_step'],run['kc_event_node'],lookup,u,w,w/3,dt)
                total=float(edge_elig(own,source_local,comp).sum(dtype=np.float64))
                effects.append({**spec,'condition':condition,'window_ms':w,'US_step':u,'fixed_KC_credit':float(c[u]),
                  'rerun_KC_credit':total,'network_feedback_credit_difference':total-float(c[u]),
                  'KC_event_stream_exact_vs_noUS':bool(np.array_equal(run['kc_event_step'],bev) and np.array_equal(run['kc_event_node'],bn)),
                  'KC_count_L1_vs_noUS':int(np.abs(run['kc_counts']-baseline['kc_counts']).sum()),
                  'MBON_count_L1_vs_noUS':int(np.abs(run['mbon_counts']-baseline['mbon_counts']).sum())})
            trace,last=last_spike_trace(bev,bn,lookup,900,w,w/3,dt)
            for j in np.flatnonzero(comp):
                ls=int(last[source_local[j]])
                edge_rows.append({'label':spec['label'],'seed':spec['seed'],'split':spec['split'],'window_ms':w,
                  'edge_id':int(ids[j]),'source_neuron_idx':int(identity.source_neuron_idx.iloc[j]),
                  'source_body_id':int(identity.source_body_id.iloc[j]),'target_neuron_idx':int(identity.target_neuron_idx.iloc[j]),
                  'frozen_group':str(identity.identity_group.iloc[j]),'compartment':cfg['target_compartment'],
                  'last_emission_step':ls,'last_arrival_step':ls+rt.delay_steps if ls>=0 else None,
                  'age_at_US_ms':float((900-ls)*dt) if ls>=0 else None,'eligibility':float(trace[source_local[j]])})
        # Actual full-brain repeats: nine additional runs + first matched run = ten.
        if ix==0:
            ref=runs['matched_90ms']
            for rep in range(1,cfg['runtime_repeats']):
                rr=sim.run(targets,rates,spec['actual_seed'],cfg['duration_ms'],dan_targets=dan,dan_pulses=[90],collect_events=True)
                exact={k:bool(np.array_equal(rr[k],ref[k])) for k in ('kc_event_step','kc_event_node','counts','final_v','final_g')}
                repeat_rows.append({'repeat':rep+1,'exact':exact,'all_exact':all(exact.values()),'wall_seconds':rr['wall_seconds']})
                assert all(exact.values())
        print(json.dumps({'completed_trials':ix+1,'of':len(inputs),'label':spec['label'],'seed':spec['seed']}),flush=True)
    for name,rows in [('events/trials',telemetry),('timing/KC-histograms',hist),('fixed-event/ratios',fixed),
                      ('fixed-event/time-grid',grid),('network-effects/credit-decomposition',effects),('timing/edge-age-at-90ms',edge_rows)]:
        pd.DataFrame(rows).to_parquet(exp/(name+'.parquet'),index=False)
    save(exp/'regression/phase4q-reproduction.json',{'rows':q_repro,'all_event_and_credit_exact':all(x['event_hash_exact'] and x['credit_exact'] for x in q_repro)})
    save(exp/'regression/independent-oracle.json',{'rows':oracle_rows,'all_exact':all(x['exact'] for x in oracle_rows)})
    save(exp/'regression/full-runtime-10-repeat.json',{'scope':'one calibration odor/seed, 120ms, unchanged APL and runtime-v2',
         'runs':10,'rows':repeat_rows,'all_exact':all(x['all_exact'] for x in repeat_rows)})
    # Runtime weights are restored before comparing full frozen five-trial references.
    rt.weights.copy_(sim.base);reg=[]
    for spec in pcfg['reference_trials']:
        refroot=Path(cfg['phase4p_root'])/'phase4c-v2/reference'/spec['name']
        ref_events=pd.read_parquet(refroot/'spike-events.parquet')
        ref={'step':ref_events.step.to_numpy(),'neuron_idx':ref_events.neuron_idx.to_numpy(),
          'counts':np.load(refroot/'per-neuron-spike-counts.int32.npy'),'v':np.load(refroot/'final-voltage.float32.npy'),'g':np.load(refroot/'final-g.float32.npy')}
        for rep in range(cfg['reference_repeats']):
            r=rt.run(np.flatnonzero(rt.pops[spec['population']]),spec['rate_hz'],spec['seed'],spec['duration_ms'])
            exact={k:bool(np.array_equal(r[k],ref[k])) for k in ref}
            reg.append({'trial':spec['name'],'repeat':rep+1,'exact':exact,'all_exact':all(exact.values()),'wall_seconds':r['metrics']['wall_seconds']})
            assert all(exact.values()), f'Frozen regression failed: {spec}'
        print(json.dumps({'regression':spec['name'],'exact':True}),flush=True)
    save(exp/'regression/five-frozen-trials.json',{'all_exact':all(x['all_exact'] for x in reg),'rows':reg,'reference':'Phase4P phase4c-v2 frozen reference'})
    # Eligibility controls operate on actual recorded events, and are explicitly offline.
    base=runs['no_external_US'];maxstep=int(base['kc_event_step'].max()) if len(base['kc_event_step']) else 0
    ctr=[]
    for w in cfg['windows_ms']:
        before,_=last_spike_trace(base['kc_event_step'],base['kc_event_node'],lookup,-10000,w,w/3,dt)
        after,_=last_spike_trace(base['kc_event_step'],base['kc_event_node'],lookup,maxstep+10000,w,w/3,dt)
        at,last=last_spike_trace(base['kc_event_step'],base['kc_event_node'],lookup,900,w,w/3,dt)
        path=exp/'regression'/f'checkpoint-window-{w}.npz';np.savez_compressed(path,last=last)
        restored=np.load(path)['last'];age=(900-restored).astype(np.float32)*np.float32(dt)
        reconstructed=np.zeros(len(last),np.float32);ok=(restored>=0)&(age<=np.float32(w));reconstructed[ok]=np.exp(-age[ok]/np.float32(w/3))
        ctr.append({'window_ms':w,'far_before_zero':not np.any(before),'far_after_zero':not np.any(after),
           'D0_zero':not np.any(edge_elig(at,source_local,comp,False)),
           'unrelated_compartment_zero_on_scored_edges':not np.any(edge_elig(at,source_local,identity.compartment.astype(str).to_numpy()=="y2a'1")[comp]),
           'disk_checkpoint_restore_exact':bool(np.array_equal(at,reconstructed)),
           'future_append_invariant':bool(np.array_equal(at,oracle(base['kc_event_step'],base['kc_event_node'],lookup,900,w,w/3,dt)))})
    save(exp/'regression/offline-controls.json',{'scope':'actual events, offline eligibility only; not biological controls','rows':ctr})
    cpu1=resource.getrusage(resource.RUSAGE_SELF);wall=time.perf_counter()-t0
    save(exp/'performance/summary.json',{'wall_seconds_audit_and_regression':wall,'brain_runs_main':len(inputs)*4,
        'extra_full_runtime_repeats':9,'frozen_reference_runs':len(reg),'main_brain_wall_seconds':sum(walls),
        'CPU_percent_one_core':100*(cpu1.ru_utime+cpu1.ru_stime-cpu0.ru_utime-cpu0.ru_stime)/wall,
        'peak_RAM_bytes':resource.getrusage(resource.RUSAGE_SELF).ru_maxrss*1024,'peak_VRAM_bytes':torch.cuda.max_memory_allocated(),
        'GPU_name':torch.cuda.get_device_name(0),'GPU_utilization':'not sampled; not inferred from timing'})
    analyze(cfg,exp)


def analyze(cfg,exp):
    fixed=pd.read_parquet(exp/'fixed-event/ratios.parquet');ev=pd.read_parquet(exp/'events/trials.parquet')
    eff=pd.read_parquet(exp/'network-effects/credit-decomposition.parquet');age=pd.read_parquet(exp/'timing/edge-age-at-90ms.parquet')
    ratio=[]
    for (split,w),f in fixed.groupby(['split','window_ms']):
        m=float(f.matched_same_events.mean());q=float(f.random_q_exact_expectation.mean());n=float(f.random_n_exact_expectation.mean())
        ratio.append({'split':split,'window_ms':int(w),'matched':m,'random_Q_exact':q,'random_N_exact':n,
                      'ratio_Q':m/q if q else None,'ratio_N':m/n if n else None})
    groups=age.groupby(['split','window_ms','frozen_group'],as_index=False).agg(total_eligibility=('eligibility','sum'),credited_observations=('eligibility',lambda x:int((x>0).sum())))
    groups.to_parquet(exp/'audit/frozen-group-credit.parquet',index=False)
    report={'phase':'4R','status':'COMPLETED-DIAGNOSTIC-NO-CANDIDATE','learning_run':False,'next_phase_started':False,
      'new_timing_selected':None,'new_eligibility_rule_selected':None,'formal_AB_evaluated':False,
      'same_event_ratios':ratio,'symmetric_delay_exact':bool(fixed.symmetric_delay_exact.all()),
      'activity_by_condition':ev.groupby(['split','condition'],as_index=False)[['KC_events','KC_active','compartment_fraction_first_30ms','compartment_events_85_to_90ms']].mean().to_dict('records'),
      'network_effects':eff.groupby(['split','condition','window_ms'],as_index=False)[['fixed_KC_credit','rerun_KC_credit','network_feedback_credit_difference','KC_count_L1_vs_noUS','MBON_count_L1_vs_noUS']].mean().to_dict('records'),
      'limitations':['Same two calibration odors with independent heldout seeds; not heldout odors or generalization proof.',
                    'No APL ablation; early transient cannot be uniquely attributed to APL.',
                    'Emission-time vs arrival-time semantics audited without changing frozen teaching semantics.',
                    'No training or functional associative-memory gate in this audit.',
                    'Frozen-group labels belong to formal A/B; calibration credit on these groups is not calibration odor specificity.']}
    save(exp/'audit/results.json',report)
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    hist=pd.read_parquet(exp/'timing/KC-histograms.parquet');grid=pd.read_parquet(exp/'fixed-event/time-grid.parquet')
    fig,axs=plt.subplots(2,2,figsize=(13,8),layout='constrained')
    for oi,label in enumerate(('CAL_A','CAL_B')):
        for split,style in [('calibration','-'),('heldout','--')]:
            h=hist[(hist.label==label)&(hist.split==split)&(hist.condition=='no_external_US')].groupby('start_ms').compartment_source_spikes.mean()
            axs[0,oi].plot(h.index+2.5,h.values,style,label=split)
        axs[0,oi].axvline(90,color='black',ls=':',label='US 90 ms');axs[0,oi].set(title=label+' target-compartment KC activity',xlabel='Time (ms)',ylabel='KC spikes / 5 ms');axs[0,oi].legend()
        for w in cfg['windows_ms']:
            g=grid[(grid.label==label)&(grid.split=='heldout')&(grid.window_ms==w)].groupby('US_step').credit_same_KC_events.mean()
            axs[1,oi].plot(g.index*0.1,g.values,label=f'{w} ms')
        axs[1,oi].axvline(90,color='black',ls=':');axs[1,oi].set(title='Fixed-event credit vs US time (heldout seeds)',xlabel='US time (ms)',ylabel='Summed edge eligibility');axs[1,oi].legend(ncol=2)
    fig.savefig(exp/'figures/timing-audit.png',dpi=180);plt.close(fig)
    print(json.dumps({'completed':True,'ratios':ratio}),flush=True)


def seal(cfg,exp,model):
    frozen=phase4p_finalize.integrity();after=quick_integrity(cfg)
    before=json.loads((exp/'audit/preflight.json').read_text())
    source=json.loads((exp/'protocol/source-hashes.json').read_text())
    same=all(digest(Path(p))==h for p,h in source.items())
    rep={'Phase1_4O':frozen,'Phase4P_4Q':after,'same_before_after':before==after,'source_hashes_unchanged':same,
         'all_passed':frozen['all_static_checks_passed'] and all(x['passed'] for x in after.values()) and before==after and same}
    save(exp/'audit/frozen-integrity.json',rep)
    for root in (model,exp):
        paths=[p for p in sorted(root.rglob('*')) if p.is_file() and p.name not in ('checksums.sha256','manifest.json')]
        rows=[{'path':p.relative_to(root).as_posix(),'sha256':digest(p),'bytes':p.stat().st_size} for p in paths]
        save(root/'manifest.json',{'root':str(root),'files':rows})
        rows.append({'path':'manifest.json','sha256':digest(root/'manifest.json')})
        (root/'checksums.sha256').write_text(''.join(x['sha256']+'  '+x['path']+'\n' for x in rows))
        own=verify_checksum(root,root/'checksums.sha256');assert own['passed'];print(json.dumps(own),flush=True)
    print(json.dumps({'frozen_integrity':rep['all_passed']}),flush=True)


if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('task',choices=['audit','seal']);parser.add_argument('--config',required=True)
    args=parser.parse_args();cfg=yaml.safe_load(Path(args.config).read_text());exp=Path(cfg['experiment_root']);model=Path(cfg['model_root'])
    if args.task=='audit':run_audit(cfg,exp,model)
    else:seal(cfg,exp,model)