"""Isolated in-memory hypothesis overlays; all frozen model/data files read-only."""
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,phase4x
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):
    xc=yaml.safe_load((Path(cfg['prior_model'])/'phase4x.yaml').read_text());out=phase4x.checks(xc)
    for key in ('prior_root','prior_model'):
        root=Path(cfg[key]);out['X_'+key]=verify_checksum(root,root/'checksums.sha256')
    assert all(x['passed'] for x in out.values());return out

class CountObserver:
    def __init__(self,rt,step):self.rt=rt;self.ext=rt.ext;self.step=step;self.before_late=None
    def __getattr__(self,name):return getattr(self.ext,name)
    def state_step(self,*args):
        if self.rt.step_index==self.step:self.before_late=self.rt.counts.cpu().numpy().copy()
        return self.ext.state_step(*args)

def run(cfg,exp,model):
    for d in ('protocol','audit','events','states','observations','comparisons','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())
    oc,panel,vectors,selection,mapping=door_inputs();panel=panel[panel.split=='calibration'].reset_index(drop=True);assert len(panel)==6
    save(exp/'protocol/preregistered.json',{'config':cfg,'input_config':vc,'panel':panel.to_dict('records'),
        'sources':{p.name:digest(p) for p in model.iterdir() if p.is_file()},'before_simulation':True,
        'no_candidate_selection':True,'new_seed_is_one_replication_not_generalization_gate':True})
    pcfg=yaml.safe_load((Path(cfg['runtime_root'])/'phase4p.yaml').read_text());rt=phase4p.CanonicalRuntime(pcfg,Path(cfg['runtime_root']),exp);sim=DeterministicAPLSim(rt)
    role=roles(sim.meta);dpm=np.flatnonzero(role.annotation_DPM).astype(np.int32)
    assert sorted(sim.meta.iloc[dpm].body_id.astype(int).tolist())==cfg['perturbation']['expected_bodies']
    ids=np.flatnonzero(np.isin(sim.sources,dpm)).astype(np.int64);eg=torch.as_tensor(ids,dtype=torch.int64,device=sim.device)
    assert not set(ids)&set(sim.apl_out_ids)
    sim.meta.iloc[dpm].to_parquet(exp/'audit/DPM-source-metadata.parquet',index=False)
    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],'frozen_effective_weight':sim.base[eg].cpu().numpy()}).to_parquet(exp/'audit/perturbed-existing-edges.parquet',index=False)
    original_install=sim.install;original_base_hash=arrhash(sim.base.cpu().numpy());weightchecks=[]
    original_install(None);original_overlay=rt.weights.cpu().numpy().copy()
    factor=1.
    def install(E):
        assert E is None,'Learning forbidden in diagnostic branch'
        original_install(None)
        rt.weights[eg]=sim.base[eg]*float(factor)
    sim.install=install
    for spec in cfg['conditions']:
        factor=spec['factor'];sim.install(None);got=rt.weights.cpu().numpy()
        expected=original_overlay.copy();expected[ids]=original_overlay[ids]*np.float32(factor)
        assert np.array_equal(got,expected)
        changed=np.flatnonzero(got!=original_overlay)
        assert np.all(np.isin(changed,ids))
        weightchecks.append({'condition':spec['name'],'factor':factor,'changed_effective_edges':len(changed),
            'only_declared_DPM_edges':True,'edge_array_length':len(got),'raw_weight_changes':0})
    save(exp/'regression/overlay-scope.json',weightchecks)
    populations={'all':np.arange(sim.n),'KC':sim.kc,'PN_annotation':sim.pn,'MBON':sim.mbon,'DPM':dpm,
        'DAN_annotation':np.flatnonzero(role.annotation_DAN),'MBON_downstream':sim.downstream}
    for cls in ('descending','motor'):
        populations[cls]=np.flatnonzero(sim.meta.annotation_superclass.eq(cls).to_numpy())
    save(exp/'audit/populations.json',{k:{'count':len(v),'definition':('exact annotation_superclass' if k in ('descending','motor') else 'frozen IDs or exact annotation evidence')} for k,v in populations.items()})
    target=np.unique(sim.identity.loc[sim.identity.compartment.eq(vc['target_compartment']),'source_neuron_idx'].to_numpy(np.int32))
    reg=[];stats=[];watchrows=[];poprows=[];comparisons=[];inputs=[]
    start=time.perf_counter();torch.cuda.reset_peak_memory_stats()
    def trial(targets,forced,name):
        observer=CountObserver(rt,round(cfg['late_start_ms']/sim.dt));rt.ext=observer
        try:r=phase4v.observed(sim,targets,forced,target,vc)
        finally:rt.ext=observer.ext
        assert r['finite'],'Nonfinite state; stop branch analysis'
        r['late_counts']=r['counts']-observer.before_late
        assert np.all(r['late_counts']>=0)
        return r
    for oi in cfg['odor_indices']:
        odor=panel.iloc[oi];targets,rates=make_drive(oc,panel,vectors,selection,mapping,odor.odor_id,'physiological')
        for seed in cfg['seeds']:
            actual=vc['seed_namespace']+oi*vc['odor_seed_offset']+seed
            forced=phase4v.forcing(sim,targets,rates,actual,vc,'offset')
            inputs.append({'odor':oi,'seed':seed,'actual_seed':actual,'forcing_hash':arrhash(forced.cpu().numpy()),'targets_rates_hash':arrhash(targets,rates)})
            reference=None
            for spec in cfg['conditions']:
                factor=spec['factor'];condition=spec['name'];name=f'{oi}-{seed}-{condition}'
                r=trial(targets,forced,name)
                phase4s.persist_run(exp,name,r)
                np.savez_compressed(exp/'states'/f'{name}.npz',**{k:r[k] for k in ('counts','late_counts','final_v','final_g','trace')})
                pd.DataFrame(r['samples']).to_parquet(exp/'observations'/f'{name}.parquet',index=False)
                if condition=='frozen_reference':
                    reference=r
                    if seed==cfg['old_replay_seed']:
                        oldname=f'{oi}-{seed}-offset';old=np.load(Path(cfg['input_experiment'])/'states'/f'{oldname}.npz')
                        eq={k:bool(np.array_equal(old[k],r[k])) for k in ('counts','final_v','final_g','trace')}
                        ev=pd.read_parquet(Path(cfg['input_experiment'])/'events'/f'{oldname}.parquet')
                        eq['KC_events']=arrhash(ev.step.to_numpy(),ev.neuron_idx.to_numpy())==arrhash(r['kc_event_step'],r['kc_event_node'])
                        assert all(eq.values());reg.append({'trial':name,'test':'frozen_V_reproduction',**eq})
                assert reference is not None
                if oi==cfg['odor_indices'][0]:
                    rr=trial(targets,forced,name)
                    keys=('counts','late_counts','final_v','final_g','kc_event_step','kc_event_node')
                    eq={k:bool(np.array_equal(r[k],rr[k])) for k in keys};assert all(eq.values())
                    reg.append({'trial':name,'test':'exact_repeat',**eq})
                es=r['kc_event_step'];en=r['kc_event_node'];late=es>=round(cfg['late_start_ms']/sim.dt)
                stats.append({'odor_index':oi,'odor_name':odor.odor_name,'seed':seed,'condition':condition,
                    'whole_brain_spikes':int(r['counts'].astype(np.int64).sum()),'whole_brain_active':int(np.count_nonzero(r['counts'])),
                    'late_whole_brain_spikes':int(r['late_counts'].astype(np.int64).sum()),'KC_spikes':len(es),'late_KC_spikes':int(late.sum()),
                    'late_KC_active':int(np.unique(en[late]).size),'finite':r['finite'],'wall_seconds':r['wall_seconds']})
                for node in cfg['watch_neuron_idx']:
                    t=es[en==node];watchrows.append({'odor_index':oi,'seed':seed,'condition':condition,'neuron_idx':node,
                        'total_spikes':len(t),'late_spikes':int((t>=round(cfg['late_start_ms']/sim.dt)).sum()),
                        'first_ms':float(t.min()*sim.dt) if len(t) else None,'last_ms':float(t.max()*sim.dt) if len(t) else None})
                for pop,idx in populations.items():
                    poprows.append({'odor_index':oi,'seed':seed,'condition':condition,'population':pop,'neurons':len(idx),
                        'total_spikes':int(r['counts'][idx].astype(np.int64).sum()),'active':int(np.count_nonzero(r['counts'][idx])),
                        'late_spikes':int(r['late_counts'][idx].astype(np.int64).sum()),'late_active':int(np.count_nonzero(r['late_counts'][idx]))})
                comparisons.append({'odor_index':oi,'seed':seed,'condition':condition,
                    'neurons_with_changed_counts':int(np.count_nonzero(r['counts']!=reference['counts'])),
                    'max_abs_finalV_difference':float(np.abs(r['final_v'].astype(float)-reference['final_v']).max()),
                    'KC_event_sequence_equal':arrhash(es,en)==arrhash(reference['kc_event_step'],reference['kc_event_node'])})
            save(exp/'regression/trials.json',reg)
            print(json.dumps({'odor_complete':oi,'seed':seed,'trials':len(stats)}),flush=True)
    sim.install=original_install;original_install(None)
    assert arrhash(sim.base.cpu().numpy())==original_base_hash
    assert np.array_equal(rt.weights.cpu().numpy(),original_overlay)
    save(exp/'regression/restore.json',{'immutable_memory_base_hash_exact':True,'frozen_APL_overlay_restored_exact':True})
    assert phase4s.evidence_regression(cfg,exp,rt,sim)
    for name,data in [('trials',stats),('watched-KC',watchrows),('populations',poprows),('paired-differences',comparisons)]:pd.DataFrame(data).to_parquet(exp/'comparisons'/f'{name}.parquet',index=False)
    save(exp/'protocol/input-hashes.json',inputs)
    save(exp/'audit/results.json',{'phase':'4Y','status':'COMPLETE-EXPERIMENTAL-HYPOTHESIS-AUDIT','trials':len(stats),'repeat_trials':6,
        'old_reference_reproductions':6,'frozen_reference_trials':5,'reference_pass':True,'all_finite':True,
        'DPM_neurons':len(dpm),'DPM_existing_edges':len(ids),'weight_training':False,'candidate_selected':False,
        'frozen_baseline_changed':False,'biological_sign_proven':False,'new_seed_count':1,'old_seed_count':1})
    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_observation_overhead_included':True})

def seal(cfg,exp,model):
    pre=json.loads((exp/'protocol/preregistered.json').read_text());now=checks(cfg);hist=phase4p_finalize.integrity()
    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['sources'].items())
    save(exp/'audit/frozen-integrity.json',{'pass':ok,'history':hist,'current':now});assert ok
    result=json.loads((exp/'audit/results.json').read_text());result['integrity_pass']=True;save(exp/'audit/results.json',result)
    w=pd.read_parquet(exp/'comparisons/watched-KC.parquet');t=pd.read_parquet(exp/'comparisons/trials.parquet')
    lines=['# Phase 4Y — independent DPM fast-current hypotheses','',
        'Status: COMPLETE-EXPERIMENTAL-HYPOTHESIS-AUDIT. No candidate selection, training, game integration or frozen baseline changes.',
        'All conditions use identical PN forcing for 120ms, then off to 1000ms. Six old odors, one old seed and one new seed. Not a population-level independent generalization gate.',
        'Factors 0 and -1 are diagnostic model interventions, not corrections to raw labels or validated receptor-specific physiology. Only actual outgoing DPM effective currents change in RAM. Graph and raw counts stay intact.','',
        '| KC | condition | late spike range (12 trials) |','|---:|---|---:|']
    for (node,condition),g in w.groupby(['neuron_idx','condition']):lines.append(f'| {node} | {condition} | {g.late_spikes.min()}–{g.late_spikes.max()} |')
    lines+=['','| condition | mean whole-brain spikes | mean late KC spikes | mean late active KC |','|---|---:|---:|---:|']
    for condition,g in t.groupby('condition'):lines.append(f'| {condition} | {g.whole_brain_spikes.mean():.2f} | {g.late_KC_spikes.mean():.2f} | {g.late_KC_active.mean():.2f} |')
    lines+=['','Exact controls: six Phase4V original trials, six repeats, five frozen reference trials, unchanged in-memory base hash and restored frozen overlay.',
        'Finite states do not establish biological stability. No calibrated runaway threshold, biological sign conclusion, receptor kinetics or learning success is inferred.',
        'Whole-brain intervention outcomes may differ from prior isolated replay because all feedback responds. Report this difference, do not select a condition by favorable activity.']
    (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));q=verify_checksum(root,root/'checksums.sha256');assert q['passed'];print(json.dumps(q),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']))