"""Event-resolved read-only diagnosis of frozen trials and existing 4Y hypotheses."""
from pathlib import Path
import argparse,json,time
import numpy as np
import pandas as pd
import torch,yaml
import phase4p,phase4p_finalize,phase4s,phase4v,phase4w,phase4y
from phase4p_revalidate import DeterministicAPLSim,door_inputs,make_drive
from phase4r import save,digest,arrhash
from phase4h_audit import verify_checksum
from transmitter_roles_v1 import roles


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


class Observer(phase4w.Observer):
    def __init__(self,rt,watch,requests,targets):
        super().__init__(rt,watch)
        self.requests=requests;self.targets=np.asarray(targets,np.int32);self.request_rows=[]
    def force_step(self,*a):
        step=self.rt.step_index;ii=np.flatnonzero(self.requests[step])
        if len(ii):
            tg=a[4][torch.as_tensor(ii,device=self.rt.device)]
            x=torch.stack((a[0][tg],self.rt.g[tg],a[1][tg].float(),a[2][tg].float())).cpu().numpy()
        self.ext.force_step(*a)
        if len(ii):
            after=a[2][tg].cpu().numpy()
            for j,k in enumerate(ii):
                node=int(self.targets[k]);status='preexisting' if x[3,j] else ('blocked' if x[2,j] else 'accepted')
                self.request_rows.append({'step':step,'neuron_idx':node,'status':status,
                    'pre_force_v':float(x[0,j]),'pre_force_g':float(x[1,j]),'pre_force_ref':int(x[2,j]),
                    'preexisting_spike':bool(x[3,j]),'actual_output_spike':bool(after[j])})
                assert bool(after[j])==(status!='blocked')
                if status=='accepted':self.forced.append((step,node))


def role_codes(meta):
    x=roles(meta);out=x.role_group.to_numpy(object).copy()
    for cls,code in [('ALPN','ALPN'),('ALLN','ALLN'),('Kenyon_Cell','KC'),('MBON','MBON'),('DAN','DAN')]:
        out[meta.annotation_class.eq(cls).to_numpy()]=code
    out[meta.annotation_type.eq('DPM').to_numpy()]='DPM'
    out[meta.annotation_type.eq('APL').to_numpy()]='APL'
    return out


def replay(arrival,overlay,requests,rt):
    f=np.float32;v=f(rt.rest);g=f(0);ref=0;events=[];states=[];accepted=[]
    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
        states.append((v,g,float(fire)))
        if requests[step] and not fire and ref==0:
            v=f(rt.resetv);ref=rt.ref_steps;fire=True;accepted.append(step)
        if fire:events.append(step)
        g=f(np.float64(g)+overlay[step])
    return np.asarray(events,np.int32),np.asarray(states,np.float32),v,g,np.asarray(accepted,np.int32)


def summarize_drive(cfg,exp,name,sim,obs,result,edges,codes,request_frame):
    rt=sim.rt;state=np.asarray(obs.states);T=len(state);ev=obs.events()
    assert np.array_equal(np.bincount(ev.neuron_idx,minlength=sim.n),result['counts'])
    arrivals=ev.merge(edges,left_on='neuron_idx',right_on='source_neuron_idx').drop(columns='neuron_idx')
    arrivals=arrivals.rename(columns={'step':'emission_step'})
    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')
    # Compact event-derived edge accounting retains exact edge IDs; whole events permit full reconstruction.
    summaries=[];replays=[];reg=[]
    for wi,node in enumerate(obs.watch):
        a=arrivals[arrivals.target_neuron_idx==node]
        grouped={int(t):g for t,g in a.groupby('arrival_step',sort=False)}
        ar=np.zeros(T,np.float32)
        for t,g in grouped.items():ar[t]=np.cumsum(g.effective_weight.to_numpy(np.float32),dtype=np.float32)[-1]
        assert np.array_equal(ar,state[:,3,wi]),f'arrival mismatch {node}'
        after=np.r_[state[1:,1,wi],result['final_g'][node]]
        overlay=after.astype(np.float64)-state[:,5,wi].astype(np.float64)
        if codes[node]!='KC':assert not np.any(overlay),'unexpected non-KC overlay'
        req=np.zeros(T,bool);q=request_frame[request_frame.neuron_idx==node];req[q.step.to_numpy(np.int32)]=True
        base=replay(ar,overlay,req,rt)
        expected=ev[ev.neuron_idx==node].step.to_numpy(np.int32)
        eq={'queue_exact':True,'events_exact':np.array_equal(base[0],expected),
            'post_state_exact':np.array_equal(base[1],state[:,4:7,wi]),
            'final_v_exact':base[2]==result['final_v'][node],'final_g_exact':base[3]==result['final_g'][node],
            'accepted_exact':np.array_equal(base[4],q[q.status=='accepted'].step.to_numpy(np.int32))}
        assert all(eq.values()),(name,int(node),eq)
        reg.append({'trial':name,'neuron_idx':int(node),**{k:bool(v) for k,v in eq.items()}})
        # Predeclared population omissions plus each actual incoming ALPN edge to the target KC.
        modes=[('base',None)]+[('omit_'+c,a.source_role.eq(c).to_numpy()) for c in cfg['isolated_omissions']]
        if node==cfg['target_KC']:
            modes += [('omit_edge_'+str(e),a.edge_id.eq(e).to_numpy()) for e in edges.loc[(edges.target_neuron_idx==node)&(edges.source_role=='ALPN'),'edge_id']]
        for mode,omit in modes:
            aa=ar
            if omit is not None:
                aa=np.zeros(T,np.float32)
                for t,g in a.loc[~omit].groupby('arrival_step',sort=False):
                    aa[int(t)]=np.cumsum(g.effective_weight.to_numpy(np.float32),dtype=np.float32)[-1]
            rr=base if omit is None else replay(aa,overlay,req,rt)
            replays.append({'trial':name,'neuron_idx':int(node),'mode':mode,'spikes':len(rr[0]),
                'late_spikes':int(np.count_nonzero(rr[0]>=round(cfg['late_start_ms']/rt.dt))),
                'accepted_requests':len(rr[4]),'max_abs_v_diff':float(np.abs(rr[1][:,0].astype(float)-base[1][:,0]).max())})
        for period,lo,hi in [('input_on',0,1200),('post_input',1200,T),('late',5000,T)]:
            aa=a[(a.arrival_step>=lo)&(a.arrival_step<hi)]
            for eid,g in aa.groupby('edge_id',sort=False):
                e=g.iloc[0];weight=float(e.effective_weight)
                summaries.append({'trial':name,'period':period,'target_neuron_idx':int(node),'edge_id':int(eid),
                    'source_neuron_idx':int(e.source_neuron_idx),'source_role':str(e.source_role),
                    'raw_synapse_count':int(e.raw_synapse_count),'effective_weight':weight,'edge_event_count':len(g),
                    'signed_arrival_sum':len(g)*weight,'positive_arrival_sum':len(g)*max(0.,weight),
                    'negative_arrival_sum':len(g)*min(0.,weight)})
    return ev,reg,replays,summaries


def request_comparison(exp,oi,seed,a,b):
    z=a.merge(b,on=['step','neuron_idx'],suffixes=('_reference','_off'),validate='one_to_one',how='outer',indicator=True)
    assert z['_merge'].eq('both').all();z=z.drop(columns='_merge')
    z['status_changed']=z.status_reference!=z.status_off
    z['state_changed']=(z.pre_force_v_reference!=z.pre_force_v_off)|(z.pre_force_g_reference!=z.pre_force_g_off)|(z.pre_force_ref_reference!=z.pre_force_ref_off)
    z.to_parquet(exp/'input-audit'/f'{oi}-{seed}-paired-requests.parquet',index=False)
    diff=z[z.status_changed];st=z[z.state_changed]
    row={'odor_index':oi,'seed':seed,'scheduled_requests':len(z),'status_changed_requests':len(diff),
        'first_status_difference_ms':float(diff.step.min()/10) if len(diff) else None,
        'first_requested_state_difference_ms':float(st.step.min()/10) if len(st) else None,
        'delivered_output_differences':int((z.actual_output_spike_reference!=z.actual_output_spike_off).sum()),
        'neurons_with_status_difference':int(diff.neuron_idx.nunique()),'requests_exact':True}
    for cond,q in [('reference',a),('off',b)]:
        for status in ('accepted','preexisting','blocked'):row[cond+'_'+status]=int(q.status.eq(status).sum())
    return row


def run(cfg,exp,model):
    for d in ('protocol','audit','events','states','drive','replay','ancestry','input-audit','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['input_model'])/'phase4v.yaml').read_text())
    assert vc['short_duration_ms']==120 and vc['observation_duration_ms']==1000
    save(exp/'protocol/preregistered.json',{'config':cfg,'source_hashes':{p.name:digest(p) for p in model.iterdir() if p.is_file()},
        'before_simulation':True,'selection_is_exploratory':True,'input_request_clamp_not_attempted':True})
    pc=yaml.safe_load((Path(cfg['runtime_root'])/'phase4p.yaml').read_text())
    rt=phase4p.CanonicalRuntime(pc,Path(cfg['runtime_root']),exp);sim=DeterministicAPLSim(rt);sim.install(None)
    assert sim.n==166700 and len(rt.raw_np)==25582938 and rt.dt==0.1 and rt.delay_steps==18
    codes=role_codes(sim.meta);original_install=sim.install;basehash=arrhash(sim.base.cpu().numpy());weights0=rt.weights.cpu().numpy().copy()
    target=cfg['target_KC'];incoming=np.flatnonzero((rt.indices_np==target)&(codes[sim.sources]=='ALPN'))
    pn=np.unique(sim.sources[incoming]);watch=np.unique(np.r_[target,pn,cfg['local_loop_nodes']]).astype(np.int32)
    assert codes[target]=='KC' and all(codes[cfg['local_loop_nodes']]=='ALLN')
    sim.meta.iloc[watch].to_parquet(exp/'audit/watched-metadata.parquet',index=False)
    save(exp/'protocol/watched.json',{'target_KC':target,'all_incoming_ALPNs':pn.tolist(),'ALLN_from_prior_event_path':cfg['local_loop_nodes'],'all_watched':watch.tolist()})
    dpm=np.flatnonzero(sim.meta.annotation_type.eq('DPM'));de=np.flatnonzero(np.isin(sim.sources,dpm));dg=torch.as_tensor(de,device=sim.device)
    assert len(dpm)==2 and len(de)==4934
    factor=1.
    def install(E):
        assert E is None;original_install(None);rt.weights[dg]=sim.base[dg]*float(factor)
    sim.install=install
    ids=np.flatnonzero(np.isin(rt.indices_np,watch))
    edge_base=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],'source_role':codes[sim.sources[ids]]})
    edge_base.assign(frozen_effective_weight=weights0[ids]).to_parquet(exp/'audit/incoming-edges.parquet',index=False)
    sim.meta.iloc[np.unique(sim.sources[ids])].to_parquet(exp/'audit/incoming-source-metadata.parquet',index=False)
    csc=(np.argsort(rt.indices_np,kind='stable'),np.r_[0,np.cumsum(np.bincount(rt.indices_np,minlength=sim.n))])
    oldcode=phase4w.classify(sim.meta)
    oc,panel,vectors,selection,mapping=door_inputs();panel=panel[panel.split=='calibration'].reset_index(drop=True)
    targetkc=np.unique(sim.identity.loc[sim.identity.compartment.eq(vc['target_compartment']),'source_neuron_idx'].to_numpy(np.int32))
    regs=[];replays=[];drives=[];paths=[];pathsum=[];requests=[];trials=[];inputhash=[]
    start=time.perf_counter();torch.cuda.reset_peak_memory_stats()
    for oi in cfg['odor_indices']:
        seed=cfg['seed'];odor=panel.iloc[oi];targets,rates=make_drive(oc,panel,vectors,selection,mapping,odor.odor_id,'physiological')
        forced=phase4v.forcing(sim,targets,rates,vc['seed_namespace']+oi*vc['odor_seed_offset']+seed,vc,'offset')
        ff=forced.cpu().numpy();inputhash.append({'odor_index':oi,'forcing':arrhash(ff),'targets_rates':arrhash(targets,rates)})
        reqpair={}
        for spec in cfg['conditions']:
            factor=spec['factor'];condition=spec['name'];name=f'{oi}-{seed}-{condition}'
            obs=Observer(rt,watch,ff,targets);obs.attach()
            try:r=phase4v.observed(sim,targets,forced,targetkc,vc)
            finally:obs.detach()
            old=np.load(Path(cfg['prior_root'])/'states'/f'{name}.npz')
            eq={k:bool(np.array_equal(old[k],r[k])) for k in ('counts','final_v','final_g','trace')}
            oldev=pd.read_parquet(Path(cfg['prior_root'])/'events'/f'{name}.parquet')
            eq['KC_events']=arrhash(oldev.step.to_numpy(),oldev.neuron_idx.to_numpy())==arrhash(r['kc_event_step'],r['kc_event_node'])
            eq['finite']=r['finite'];assert all(eq.values()),(name,eq)
            regs.append({'trial':name,'prior_Y_exact':eq})
            rq=pd.DataFrame(obs.request_rows);rq.to_parquet(exp/'input-audit'/f'{name}-requests.parquet',index=False);reqpair[condition]=rq
            last=r['samples'][-1]
            assert len(rq)==int(ff.sum())==last['PN_requested_cumulative']
            for status,field in [('accepted','PN_accepted_cumulative'),('preexisting','PN_preexisting_cumulative'),('blocked','PN_blocked_cumulative')]:
                assert int(rq.status.eq(status).sum())==last[field]
            weights=rt.weights.cpu().numpy().copy();expected=weights0.copy();expected[de]*=np.float32(factor);assert np.array_equal(weights,expected)
            edges=edge_base.assign(effective_weight=weights[ids])
            ev,reg,rep,drive=summarize_drive(cfg,exp,name,sim,obs,r,edges,codes,rq)
            regs[-1]['CPU_shadow_replay']=reg;replays+=rep;drives+=drive
            ev.to_parquet(exp/'events'/f'{name}-whole-brain.parquet',index=False)
            np.savez_compressed(exp/'states'/f'{name}-watched.npz',state=np.asarray(obs.states),watch=watch)
            if condition=='frozen_reference':
                oldfull=pd.read_parquet(Path(cfg['event_experiment'])/'events'/f'{oi}-{seed}-offset-whole-brain.parquet')
                regs[-1]['W_full_events_exact']=bool(ev.equals(oldfull));assert regs[-1]['W_full_events_exact']
            ancestry_cfg={'watch_neuron_idx':[target],'late_start_ms':cfg['late_start_ms'],'ancestry':cfg['ancestry']}
            pr,ps=phase4w.ancestry(ancestry_cfg,sim,ev,obs.forced,csc,weights,oldcode,name)
            # Correct orthogonal role labels; historical W category field never used as biological DAN identity.
            for row in pr:
                row['source_category']=str(codes[row['source_neuron_idx']]);row['target_category']=str(codes[row['target_neuron_idx']])
            for item in ps:
                selected=[p for p in pr if p['endpoint']==item['endpoint']]
                item['categories']=sorted({p['source_category'] for p in selected})
                item['repeated_neuron_ids']=len({p['source_neuron_idx'] for p in selected})<len(selected)
            paths+=pr;pathsum+=ps
            trials.append({'trial':name,'odor_index':oi,'seed':seed,'condition':condition,'whole_spikes':int(r['counts'].astype(np.int64).sum()),
                'KC82312_late_spikes':int(((ev.neuron_idx==target)&(ev.step>=5000)).sum()),'finite':r['finite'],'wall_seconds':r['wall_seconds']})
            if oi==cfg['repeat_odor_index']:
                again=Observer(rt,watch,ff,targets);again.attach()
                try:rr=phase4v.observed(sim,targets,forced,targetkc,vc)
                finally:again.detach()
                repeat={'events':ev.equals(again.events()),'states':np.array_equal(obs.states,again.states),'requests':rq.equals(pd.DataFrame(again.request_rows))}
                assert all(repeat.values());regs[-1]['repeat_exact']=repeat
            save(exp/'regression/trials.json',regs)
            print(json.dumps({'trial_complete':name,'watched_neurons':len(watch),'incoming_edges':len(ids),'CPU_checks':len(reg)}),flush=True)
        requests.append(request_comparison(exp,oi,seed,reqpair['frozen_reference'],reqpair['DPM_fast_current_off']))
    sim.install=original_install;original_install(None)
    assert arrhash(sim.base.cpu().numpy())==basehash and np.array_equal(rt.weights.cpu().numpy(),weights0)
    save(exp/'regression/restore.json',{'base_unchanged':True,'APL_overlay_restored':True,'new_weight_hypotheses':False})
    assert phase4s.evidence_regression(cfg,exp,rt,sim)
    for folder,name,rows in [('audit','trials',trials),('replay','results',replays),('drive','edge-accounting',drives),('ancestry','event-supported-edges',paths),('input-audit','paired-summary',requests)]:
        pd.DataFrame(rows).to_parquet(exp/folder/f'{name}.parquet',index=False)
    save(exp/'ancestry/summary.json',pathsum);save(exp/'protocol/input-hashes.json',inputhash)
    save(exp/'audit/results.json',{'phase':'4Z','status':'COMPLETE-EVENT-RESOLVED-PN-AUDIT','trials':len(trials),'exact_repeat_trials':2,
        'Y_reproductions':len(trials),'W_full_event_reproductions':6,'CPU_shadow_replay_checks':sum(len(x['CPU_shadow_replay']) for x in regs),
        'reference_pass':True,'all_finite':True,'input_pairs':len(requests),'pairs_with_status_difference':sum(x['status_changed_requests']>0 for x in requests),
        'full_brain_loop_necessity_proven':False,'biological_sign_proven':False,'weight_training':False,'candidate':False,'frozen_baseline_changed':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),'diagnostic_overhead_included':True})


def describe(cfg,exp):
    d=pd.read_parquet(exp/'drive/edge-accounting.parquet');rep=pd.read_parquet(exp/'replay/results.parquet')
    q=pd.read_parquet(exp/'input-audit/paired-summary.parquet');meta=pd.read_parquet(exp/'audit/incoming-source-metadata.parquet')
    late=d[d.period=='late'].copy();den=late.groupby(['trial','target_neuron_idx']).positive_arrival_sum.transform('sum')
    late['positive_fraction']=np.where(den>0,late.positive_arrival_sum/den,np.nan)
    late=late.merge(meta[['neuron_idx','body_id','annotation_type','annotation_class','neurotransmitter','nt_ground_truth']],left_on='source_neuron_idx',right_on='neuron_idx',how='left',validate='many_to_one').drop(columns='neuron_idx')
    late.to_parquet(exp/'drive/late-input-fractions.parquet',index=False)
    lines=['# Phase 4Z — PN feedback and input acceptance audit','',
        'Diagnostic only. No training, no new sign hypothesis, no candidate. The two branches reproduce sealed Phase4Y reference and DPM-off conditions.',
        'Target and three ALLNs were selected from prior data; this is mechanism exploration, not heldout validation.',
        'All source fractions are positive arrival-weight accounting, not causal percentages or voltage attribution.','',
        '## Late positive inputs to KC 82312','', '| condition | source idx | body_id | type | mean fraction |','|---|---:|---:|---|---:|']
    z=late[late.target_neuron_idx==cfg['target_KC']].copy();z['condition']=z.trial.str.split('-',n=2).str[-1]
    for (cond,idx,body,typ),g in z.groupby(['condition','source_neuron_idx','body_id','annotation_type'],dropna=False):
        if g.positive_arrival_sum.sum()>0:lines.append(f'| {cond} | {idx} | {body} | {typ} | {g.positive_fraction.mean():.6f} |')
    lines+=['','## Conditional single-cell replay','',
        'Other recorded arrivals and measured APL remain fixed; original force requests and refractory logic retained. This does not test whole-brain loop necessity.',
        '| neuron_idx | condition | mode | late spike range |','|---:|---|---|---:|']
    rep['condition']=rep.trial.str.split('-',n=2).str[-1]
    for (node,cond,mode),g in rep.groupby(['neuron_idx','condition','mode']):
        if mode=='base' or node==cfg['target_KC'] or (node in cfg['local_loop_nodes'] and mode=='omit_ALLN'):
            lines.append(f'| {node} | {cond} | {mode} | {g.late_spikes.min()}–{g.late_spikes.max()} |')
    lines+=['','## Input acceptance','', '| odor | requests | accepted reference | accepted off | status differences | first difference ms |','|---:|---:|---:|---:|---:|---:|']
    for r in q.itertuples():lines.append(f'| {r.odor_index} | {r.scheduled_requests} | {r.reference_accepted} | {r.off_accepted} | {r.status_changed_requests} | {r.first_status_difference_ms} |')
    lines+=['','Identical requests do not establish matched delivered inputs. Refusal and preexisting spikes depend on endogenous state. No forcing clamp or altered refractory rule was introduced.',
        'Ancestry edges require actual emission +18 steps <= downstream spike; greedy ranking is not a unique or sufficient causal path. Repeated neuron IDs indicate a time-supported recurrent motif, not proof it is the only sustaining loop.',
        'Frozen LIF/APL, raw graph/counts, eligibility and masks unchanged. All detailed state/event checks are in regression/.']
    (exp/'REPORT.md').write_text('\n'.join(lines)+'\n',encoding='utf-8')


def seal(cfg,exp,model):
    assert not (exp/'checksums.sha256').exists(),'Already sealed'
    describe(cfg,exp)
    now=checks(cfg);hist=phase4p_finalize.integrity();pre=json.loads((exp/'protocol/preregistered.json').read_text())
    ok=hist['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['source_hashes'].items())
    save(exp/'audit/frozen-integrity.json',{'pass':ok,'history':hist,'inputs':now});assert ok
    r=json.loads((exp/'audit/results.json').read_text());r['integrity_pass']=True;save(exp/'audit/results.json',r)
    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']))