"""Side-channel event/drive audit. Frozen run function and CUDA math unchanged."""
from pathlib import Path
import argparse,json,time
import numpy as np
import pandas as pd
import torch
import yaml
import phase4p,phase4p_finalize,phase4s,phase4v
from phase4p_revalidate import DeterministicAPLSim,door_inputs,make_drive
from phase4r import save,digest,arrhash
from phase4h_audit import verify_checksum

CATS=['PN_olfactory','KC','APL','MBON','DAN','other']

def checks(cfg):
    vc=yaml.safe_load((Path(cfg['prior_model'])/'phase4v.yaml').read_text())
    out=phase4v.checks(vc)
    for key in ('prior_root','prior_model'):
        p=Path(cfg[key]);out['V_'+key]=verify_checksum(p,p/'checksums.sha256')
    assert all(x['passed'] for x in out.values());return out

class Observer:
    def __init__(self,rt,watch):
        self.rt=rt;self.ext=rt.ext;self.fanout=rt.canonical_fanout
        self.watch=np.asarray(watch,np.int32);self.wg=torch.as_tensor(watch,dtype=torch.int64,device=rt.device)
        self.states=[];self.steps=[];self.nodes=[];self.forced=[]
    def __getattr__(self,name):return getattr(self.ext,name)
    def state_step(self,*a):
        w=self.wg
        pre=torch.stack((a[0][w],a[1][w],a[2][w].float(),a[3][w]))
        self.ext.state_step(*a)
        post=torch.stack((a[0][w],a[1][w],a[4][w].float()))
        self.states.append(torch.cat((pre,post)).cpu().numpy())
    def force_step(self,*a):
        tg=a[4];req=a[5].bool()
        accepted=req & ~a[2][tg].bool() & (a[1][tg]==0)
        self.ext.force_step(*a)
        x=tg[accepted].cpu().numpy().astype(np.int32)
        if len(x):self.forced.extend((self.rt.step_index,int(i)) for i in x)
    def record_fanout(self,src,dst):
        x=src.cpu().numpy().astype(np.int32)
        self.nodes.append(x);self.steps.append(np.full(len(x),self.rt.step_index,np.int32))
        return self.fanout(src,dst)
    def attach(self):self.rt.ext=self;self.rt.canonical_fanout=self.record_fanout
    def detach(self):self.rt.ext=self.ext;self.rt.canonical_fanout=self.fanout
    def events(self):
        return pd.DataFrame({'step':np.concatenate(self.steps),'neuron_idx':np.concatenate(self.nodes)})

def classify(meta):
    cls=meta.annotation_class.fillna('').astype(str);typ=meta.annotation_type.fillna('').astype(str)
    nt=meta.neurotransmitter.fillna('').astype(str).str.strip().str.lower()
    code=np.full(len(meta),5,np.int8)
    code[cls.isin(['ALPN','olfactory'])]=0;code[cls.eq('Kenyon_Cell')]=1
    code[typ.eq('APL')]=2;code[cls.eq('MBON')]=3;code[nt.eq('dopamine')]=4
    return code

def replay(arrival,overlay,rt):
    # Explicit rounding matches CUDA __fadd_rn/__fmul_rn, not fused operations.
    f=np.float32;v=f(rt.rest);g=f(0);ref=0;out=[];vs=[];gs=[]
    for step,a in enumerate(arrival):
        g=f(f(g*rt.gd)+a);ref=max(ref-1,0)
        v=f(v+f(rt.alpha*f(g-f(v-rt.rest)))) if ref==0 else f(rt.resetv)
        fire=ref==0 and v>=rt.threshold
        if fire:v=f(rt.resetv);ref=rt.ref_steps;out.append(step)
        vs.append(v);gs.append(g)
        # Measured applied APL difference, fixed exogenous signal for isolated replay.
        g=f(np.float64(g)+overlay[step])
    return np.asarray(out,np.int32),np.asarray(vs,np.float32),np.asarray(gs,np.float32),g

def drive_audit(cfg,exp,name,sim,observer,result,edges,code):
    rt=sim.rt;st=np.asarray(observer.states);T=len(st);events=observer.events()
    assert np.array_equal(np.bincount(events.neuron_idx,minlength=sim.n),result['counts'])
    events.to_parquet(exp/'events'/f'{name}-whole-brain.parquet',index=False)
    pd.DataFrame(observer.forced,columns=['step','neuron_idx']).to_parquet(exp/'events'/f'{name}-accepted-input.parquet',index=False)
    np.savez_compressed(exp/'states'/f'{name}-watched.npz',state=st,
        fields=np.asarray(['pre_v','pre_g','pre_ref','arrival','post_v','post_g','natural_spike']))
    # Only real events on actual incoming edges. Delay is exact, including inhibitory edges.
    arrivals=events.merge(edges,left_on='neuron_idx',right_on='source_neuron_idx',how='inner')
    arrivals=arrivals.rename(columns={'step':'emission_step'}).drop(columns='neuron_idx')
    arrivals['arrival_step']=arrivals.emission_step+rt.delay_steps
    arrivals=arrivals[arrivals.arrival_step<T].sort_values(['target_neuron_idx','arrival_step','edge_id'],kind='stable')
    arrivals.to_parquet(exp/'drive'/f'{name}-actual-edge-arrivals.parquet',index=False)
    replayrows=[];drives=[];checks_out=[]
    for wi,node in enumerate(cfg['watch_neuron_idx']):
        a=arrivals[arrivals.target_neuron_idx==node]
        seqs={int(t):g for t,g in a.groupby('arrival_step',sort=False)}
        exact_arr=np.zeros(T,np.float32);bycat=np.zeros((T,len(CATS)),np.float64)
        for t,g in seqs.items():
            weights=g.effective_weight.to_numpy(np.float32)
            exact_arr[t]=np.cumsum(weights,dtype=np.float32)[-1]
            for c,q in g.groupby('source_category'):
                bycat[t,CATS.index(c)]=q.effective_weight.to_numpy(np.float64).sum()
        queue_exact=bool(np.array_equal(exact_arr,st[:,3,wi]));assert queue_exact,'Arrival reconstruction mismatch'
        after_overlay=np.r_[st[1:,1,wi],result['final_g'][node]]
        overlay=after_overlay.astype(np.float64)-st[:,5,wi].astype(np.float64)
        base=replay(exact_arr,overlay,rt)
        expected=np.flatnonzero(st[:,6,wi]).astype(np.int32)
        eq={'queue_exact':queue_exact,'events_exact':np.array_equal(base[0],expected),
            'V_exact':np.array_equal(base[1],st[:,4,wi]),'g_exact':np.array_equal(base[2],st[:,5,wi]),
            'final_g_exact':base[3]==result['final_g'][node]}
        checks_out.append({'neuron_idx':node,**{k:bool(v) for k,v in eq.items()}});assert all(eq.values()),'CPU replay divergence'
        shadow=np.zeros(len(CATS)+1,np.float64);shadowrows=[]
        for t in range(T):
            shadow*=float(rt.gd);shadow[:len(CATS)]+=bycat[t]
            row={'step':t,'neuron_idx':node,'actual_g_before_APL':float(st[t,5,wi]),
                 'shadow_residual':float(st[t,5,wi]-shadow.sum())}
            row.update({f'g_{c}':float(shadow[j]) for j,c in enumerate(CATS+['APL_overlay'])})
            shadowrows.append(row);shadow[-1]+=overlay[t]
        pd.DataFrame(shadowrows).to_parquet(exp/'drive'/f'{name}-{node}-shadow.parquet',index=False)
        for mode in cfg['isolated_replay']['modes']:
            ar=exact_arr.copy();ov=overlay.copy()
            if mode=='omit_APL_overlay':ov[:]=0
            elif mode.startswith('omit_'):
                cat=mode[5:];ar[:]=0
                for t,g in seqs.items():
                    w=g.loc[g.source_category!=cat,'effective_weight'].to_numpy(np.float32)
                    if len(w):ar[t]=np.cumsum(w,dtype=np.float32)[-1]
            rr=replay(ar,ov,rt);late=rr[0]>=round(cfg['late_start_ms']/rt.dt)
            replayrows.append({'trial':name,'neuron_idx':node,'mode':mode,'spikes':len(rr[0]),'late_spikes':int(late.sum()),
                'first_ms':float(rr[0][0]*rt.dt) if len(rr[0]) else None,'last_ms':float(rr[0][-1]*rt.dt) if len(rr[0]) else None,
                'max_abs_V_difference':float(np.abs(rr[1].astype(float)-base[1]).max())})
            np.savez_compressed(exp/'replay'/f'{name}-{node}-{mode}.npz',event_steps=rr[0],v=rr[1],g=rr[2])
        for period,minstep in [('whole',0),('late',round(cfg['late_start_ms']/rt.dt))]:
            aa=a[a.arrival_step>=minstep]
            for cat,g in aa.groupby('source_category'):
                w=g.effective_weight.to_numpy(float)
                drives.append({'trial':name,'neuron_idx':node,'period':period,'source_category':cat,
                    'positive_arrival_sum':float(w[w>0].sum()),'negative_arrival_sum':float(w[w<0].sum()),
                    'edge_event_count':len(g),'unique_sources':int(g.source_neuron_idx.nunique())})
    return checks_out,replayrows,drives,events

def ancestry(cfg,sim,events,forced,csc,weights,code,name):
    rt=sim.rt;order,ptr=csc;by=events.sort_values(['neuron_idx','step'],kind='stable')
    times=by.step.to_numpy(np.int32);ep=np.r_[0,np.cumsum(np.bincount(by.neuron_idx,minlength=sim.n))]
    roots=set(forced);rows=[];summaries=[];gd=float(rt.gd)
    def node_info(node,prefix):
        m=sim.meta.iloc[node]
        return {prefix+'body_id':int(m.body_id),prefix+'type':None if pd.isna(m.annotation_type) else str(m.annotation_type),
                prefix+'class':None if pd.isna(m.annotation_class) else str(m.annotation_class),prefix+'category':CATS[code[node]]}
    for target in cfg['watch_neuron_idx']:
        tt=times[ep[target]:ep[target+1]]
        selected=[('first',int(tt[0]))] if len(tt) else []
        late=tt[tt>=round(cfg['late_start_ms']/rt.dt)]
        if len(late):selected.append(('first_late',int(late[0])))
        for kind,t in selected:
            node=target;path=[];status='hop_limit'
            for hop in range(cfg['ancestry']['max_hops']):
                if (t,node) in roots:status='accepted_input_root';break
                ee=order[ptr[node]:ptr[node+1]];ee=ee[weights[ee]>0]
                best=None
                for e in ee:
                    src=int(sim.sources[e]);ts=times[ep[src]:ep[src+1]];j=np.searchsorted(ts,t-rt.delay_steps,side='right')-1
                    if j<0:continue
                    emit=int(ts[j]);arrival=emit+rt.delay_steps;score=float(weights[e])*gd**(t-arrival)
                    if best is None or score>best[0]:best=(score,int(e),src,emit,arrival)
                if best is None:status='unexplained_no_positive_event_predecessor';break
                score,e,src,emit,arrival=best
                row={'trial':name,'watched_neuron_idx':target,'endpoint':kind,'backward_hop':hop,
                    'source_neuron_idx':src,'target_neuron_idx':node,'source_step':emit,'arrival_step':arrival,'target_step':t,
                    'edge_id':e,'raw_synapse_count':int(rt.raw_np[e]),'effective_weight':float(weights[e]),'synaptic_sign':1,
                    'ranking_score':score,'arrival_age_ms':(t-arrival)*rt.dt,**node_info(src,'source_'),**node_info(node,'target_')}
                assert emit<t and arrival<=t
                path.append(row);node=src;t=emit
            rows.extend(path)
            summaries.append({'trial':name,'neuron_idx':target,'endpoint':kind,'hops':len(path),'status':status,
                'root_node':node,'root_step':t,'categories':sorted(set(x['source_category'] for x in path)),
                'claim':'ranked time-supported ancestry only; not unique cause or whole-brain necessity'})
    return rows,summaries

def run(cfg,exp,model):
    for d in ('protocol','audit','events','states','drive','replay','ancestry','regression','performance'):(exp/d).mkdir(parents=True,exist_ok=True)
    assert not (exp/'protocol/preregistered.json').exists(),'Refusing overwrite'
    save(exp/'audit/preflight.json',checks(cfg))
    vc=yaml.safe_load((Path(cfg['prior_model'])/'phase4v.yaml').read_text())
    save(exp/'protocol/preregistered.json',{'config':cfg,'sources':{p.name:digest(p) for p in model.iterdir() if p.is_file()},
        'before_simulation':True,'selection':'all six old odors; first fixed Phase4V seed; late common KC IDs from sealed Phase4V; diagnostic not blind validation'})
    pc=yaml.safe_load((Path(cfg['runtime_root'])/'phase4p.yaml').read_text());rt=phase4p.CanonicalRuntime(pc,Path(cfg['runtime_root']),exp)
    sim=DeterministicAPLSim(rt);code=classify(sim.meta);sim.install(None);weights=rt.weights.cpu().numpy().copy()
    assert len(sim.identity)==14232 and all(sim.meta.iloc[cfg['watch_neuron_idx']].annotation_class=='Kenyon_Cell')
    fields=['neuron_idx','body_id','annotation_class','annotation_type','annotation_subclass','annotation_superclass','neurotransmitter']
    sim.meta.iloc[cfg['watch_neuron_idx']][fields].to_parquet(exp/'audit/watched-metadata.parquet',index=False)
    ids=np.flatnonzero(np.isin(rt.indices_np,cfg['watch_neuron_idx']))
    edges=pd.DataFrame({'edge_id':ids,'source_neuron_idx':sim.sources[ids],'target_neuron_idx':rt.indices_np[ids],
        'raw_synapse_count':rt.raw_np[ids],'effective_weight':weights[ids],
        'source_category':np.asarray(CATS)[code[sim.sources[ids]]]})
    edges.to_parquet(exp/'audit/real-incoming-edges.parquet',index=False)
    sm=sim.meta.iloc[np.unique(edges.source_neuron_idx)][fields].copy();sm['audit_category']=np.asarray(CATS)[code[sm.neuron_idx]]
    sm.to_parquet(exp/'audit/incoming-source-metadata.parquet',index=False)
    save(exp/'audit/classification.json',{'PN_olfactory':'annotation_class ALPN or olfactory','KC':'annotation_class Kenyon_Cell',
        'APL':'annotation_type APL','MBON':'annotation_class MBON','DAN':'neurotransmitter dopamine (not inferred PAM/PPL1 name)',
        'other':'all remaining metadata classes; not an assertion of biological unknown',
        'APL_spiking_outgoing_weights':'frozen graded overlay zeros these; applied graded effect recorded separately'})
    csc=(np.argsort(rt.indices_np,kind='stable'),np.r_[0,np.cumsum(np.bincount(rt.indices_np,minlength=sim.n))])
    oc,panel,vectors,selection,mapping=door_inputs();panel=panel[panel.split=='calibration'].reset_index(drop=True)
    target=np.unique(sim.identity.loc[sim.identity.compartment.eq(vc['target_compartment']),'source_neuron_idx'].to_numpy(np.int32))
    reg=[];replays=[];drives=[];paths=[];pathsummary=[];start=time.perf_counter();torch.cuda.reset_peak_memory_stats()
    for oi in cfg['odor_indices']:
        seed=cfg['seed'];name=f'{oi}-{seed}-offset';odor=panel.iloc[oi]
        targets,rates=make_drive(oc,panel,vectors,selection,mapping,odor.odor_id,'physiological')
        assert not set(targets)&set(cfg['watch_neuron_idx'])
        actual=vc['seed_namespace']+oi*vc['odor_seed_offset']+seed
        forced=phase4v.forcing(sim,targets,rates,actual,vc,'offset')
        observer=Observer(rt,cfg['watch_neuron_idx']);observer.attach()
        try:r=phase4v.observed(sim,targets,forced,target,vc)
        finally:observer.detach()
        old=np.load(Path(cfg['prior_root'])/'states'/f'{name}.npz')
        oldev=pd.read_parquet(Path(cfg['prior_root'])/'events'/f'{name}.parquet')
        eq={k:bool(np.array_equal(r[k],old[k])) for k in ('counts','final_v','final_g','trace')}
        eq['KC_events']=arrhash(r['kc_event_step'],r['kc_event_node'])==arrhash(oldev.step.to_numpy(),oldev.neuron_idx.to_numpy())
        eq['finite']=r['finite'];reg.append({'trial':name,'prior_V_exact':eq});save(exp/'regression/comparisons.json',reg);assert all(eq.values())
        cs,rs,ds,events=drive_audit(cfg,exp,name,sim,observer,r,edges,code)
        reg[-1]['CPU_isolated_exact']=cs;replays.extend(rs);drives.extend(ds)
        pr,ps=ancestry(cfg,sim,events,observer.forced,csc,weights,code,name);paths.extend(pr);pathsummary.extend(ps)
        if oi==cfg['repeat_odor_index']:
            again=Observer(rt,cfg['watch_neuron_idx']);again.attach()
            try:rr=phase4v.observed(sim,targets,forced,target,vc)
            finally:again.detach()
            reg[-1]['repeat_full_events_exact']=bool(events.equals(again.events()))
            reg[-1]['repeat_watched_state_exact']=bool(np.array_equal(observer.states,again.states))
            assert reg[-1]['repeat_full_events_exact'] and reg[-1]['repeat_watched_state_exact']
        save(exp/'regression/comparisons.json',reg)
        print(json.dumps({'completed_odor':oi,'cpu_replays':len(replays),'time_supported_paths':len(pathsummary)}),flush=True)
    pd.DataFrame(replays).to_parquet(exp/'replay/results.parquet',index=False)
    pd.DataFrame(drives).to_parquet(exp/'drive/summary.parquet',index=False)
    pd.DataFrame(paths).to_parquet(exp/'ancestry/event-supported-edges.parquet',index=False)
    save(exp/'ancestry/summary.json',pathsummary)
    ref=phase4s.evidence_regression(cfg,exp,rt,sim);assert ref
    save(exp/'audit/results.json',{'phase':'4W','status':'COMPLETE-CONDITIONAL-MECHANISM-AUDIT','trials':len(cfg['odor_indices']),
        'prior_V_exact':True,'CPU_replay_exact':True,'full_event_repeat_exact':True,'reference_pass':ref,
        'rooted_time_supported_paths':sum(x['status']=='accepted_input_root' for x in pathsummary),'paths':len(pathsummary),
        'whole_brain_causal_necessity_proven':False,'learning':False,'candidate':False})
    save(exp/'performance/summary.json',{'wall_seconds_excluding_build':time.perf_counter()-start,
        'peak_VRAM_bytes':torch.cuda.max_memory_allocated(),'GPU':torch.cuda.get_device_name(0),'instrumentation_overhead_included':True})

def seal(cfg,exp,model):
    now=checks(cfg);history=phase4p_finalize.integrity();pre=json.loads((exp/'protocol/preregistered.json').read_text())
    ok=history['all_static_checks_passed'] and now==json.loads((exp/'audit/preflight.json').read_text()) and all(digest(model/n)==h for n,h in pre['sources'].items())
    save(exp/'audit/frozen-integrity.json',{'pass':ok,'history':history,'current':now});assert ok
    r=json.loads((exp/'audit/results.json').read_text());r['integrity_pass']=True;save(exp/'audit/results.json',r)
    df=pd.read_parquet(exp/'replay/results.parquet')
    lines=['# Phase 4W — late KC feedback audit','',json.dumps(r,ensure_ascii=False),'',
        'No learning or model changes. Inputs replay six old Phase4V offset trials. All frozen source artifacts remain read-only.',
        'CPU single-cell omissions hold all other recorded arrivals and the applied APL overlay fixed. They are conditional input tests, not whole-brain causal interventions.',
        'Ancestry uses real edges and actual event timing. Ranking the largest positive decayed latest-spike contribution is a heuristic, not a unique causal explanation. All inhibitory arrivals and graded APL are retained in the separate drive audit.','',
        '| neuron_idx | replay mode | mean late spikes (6 odors) |','|---:|---|---:|']
    for (node,mode),g in df.groupby(['neuron_idx','mode']):lines.append(f'| {node} | {mode} | {g.late_spikes.mean():.3f} |')
    lines += ['', 'The six reserved novel test odors remain untouched. No candidate, parameter selection, threshold/gain change, eligibility modification, training or game integration.',
        'Full evidence: drive/summary.parquet, replay/results.parquet, ancestry/event-supported-edges.parquet, regression/comparisons.json.']
    (exp/'REPORT.md').write_text('\n'.join(lines)+'\n',encoding='utf-8')
    for root in (model,exp):
        rows=[{'path':p.relative_to(root).as_posix(),'sha256':digest(p),'bytes':p.stat().st_size} for p in sorted(root.rglob('*')) if p.is_file() and p.name not in ('manifest.json','checksums.sha256')]
        save(root/'manifest.json',{'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))
        check=verify_checksum(root,root/'checksums.sha256');assert check['passed'];print(json.dumps(check),flush=True)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('task',choices=['run','seal']);p.add_argument('--config',required=True);a=p.parse_args()
    cfg=yaml.safe_load(Path(a.config).read_text());(run if a.task=='run' else seal)(cfg,Path(cfg['experiment_root']),Path(cfg['model_root']))