"""Phase 4S: independently calibrated timing, fixed-event credit, isolated transfer."""
from pathlib import Path
import argparse
import hashlib
import json
import resource
import time

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

import phase4p
import phase4p_finalize
from phase4p_revalidate import DeterministicAPLSim, IsolatedReplay, door_inputs, make_drive
from phase4h_audit import verify_checksum
from phase4q import last_spike_trace, edge_elig, update_E, group_metrics, v2_trace, replay_metrics
from phase4r import save, digest, arrhash, curves, quick_integrity


def checks(cfg):
    out=quick_integrity(cfg)
    for name,key in [('R','phase4r_root'),('R_model','phase4r_model')]:
        root=Path(cfg[key]);out[name]=verify_checksum(root,root/'checksums.sha256')
    return out


def timing_test(rows,cfg,w,u):
    x=pd.DataFrame(rows);x=x[(x.window_ms==w)&(x.US_ms==u)];details=[]
    gate=cfg['timing_gate']
    for label,g in x.groupby('label'):
        m=float(g.matched.mean());r=float(g.random_expectation.mean())
        ratio=m/r if r else None
        n=int((g.ratio>=gate['per_seed_ratio_at_least']).sum())
        details.append({'label':label,'mean_matched':m,'mean_random':r,'ratio':ratio,'qualifying_seeds':n,
          'pass':bool(ratio is not None and ratio>=gate['mean_matched_random_ratio_each_odor'] and n>=gate['qualifying_seeds_each_odor'])})
    return {'window_ms':w,'US_ms':u,'tau_ms':w/3,'odors':details,'pass':len(details)==2 and all(g['pass'] for g in details)}


def persist_run(exp,name,run):
    frame=pd.DataFrame({'step':run['kc_event_step'],'neuron_idx':run['kc_event_node']})
    frame.to_parquet(exp/'events'/f'{name}.parquet',index=False)
    return arrhash(run['kc_event_step'],run['kc_event_node'])


def evidence_regression(cfg,exp,rt,sim):
    rt.weights.copy_(sim.base);rows=[]
    for spec in rt.cfg['reference_trials']:
        ref=Path(cfg['phase4p_root'])/'phase4c-v2/reference'/spec['name'];ev=pd.read_parquet(ref/'spike-events.parquet')
        expect={'step':ev.step.to_numpy(),'neuron_idx':ev.neuron_idx.to_numpy(),'counts':np.load(ref/'per-neuron-spike-counts.int32.npy'),
          'v':np.load(ref/'final-voltage.float32.npy'),'g':np.load(ref/'final-g.float32.npy')}
        run=rt.run(np.flatnonzero(rt.pops[spec['population']]),spec['rate_hz'],spec['seed'],spec['duration_ms'])
        exact={k:arrhash(run[k])==arrhash(expect[k].astype(run[k].dtype)) for k in expect}
        rows.append({'trial':spec['name'],'exact':exact,'pass':all(exact.values())})
    save(exp/'regression/frozen-trials.json',{'rows':rows,'pass':all(x['pass'] for x in rows)})
    return all(x['pass'] for x in rows)


def run(cfg,exp,model):
    for d in ('protocol','calibration','heldout','formal','isolated','whole-brain','events','audit','regression','performance'):(exp/d).mkdir(parents=True,exist_ok=True)
    if (exp/'protocol/preregistered.json').exists():raise RuntimeError('Run already exists: refusing overwrite')
    pre=checks(cfg);save(exp/'audit/preflight.json',pre)
    assert all(v['passed'] for v in pre.values())
    ocfg,panel,vectors,selection,mapping=door_inputs();cal=panel[panel.split=='calibration'].iloc[:2]
    assert set(cal.odor_id).isdisjoint({selection['ODOR_A']['odor_id'],selection['ODOR_B']['odor_id']})
    prereg={'config':cfg,'calibration_odors':cal[['odor_id','odor_name']].to_dict('records'),
      'created_before_simulation':True,'source_hashes':{p.name:digest(p) for p in model.glob('phase4s.*')},
      'protocol_basis':'Phase4R onset response identified at 15-25ms; this informs predeclared finite grid, not formal A/B outcomes'}
    save(exp/'protocol/preregistered.json',prereg)
    pcfg=yaml.safe_load((Path(cfg['runtime_root'])/'phase4p.yaml').read_text());rt=phase4p.CanonicalRuntime(pcfg,Path(cfg['runtime_root']),exp)
    sim=DeterministicAPLSim(rt);identity=sim.identity;assert identity.identity_group.value_counts().to_dict()==cfg['formal']['fixed_groups']
    assert len(identity)==14232
    lookup=np.full(sim.n,-1,np.int32);lookup[sim.kc]=np.arange(len(sim.kc),dtype=np.int32)
    loc=lookup[identity.source_neuron_idx.to_numpy(np.int32)];assert np.all(loc>=0)
    comp=identity.compartment.astype(str).to_numpy()==cfg['target_compartment'];multi=np.bincount(loc[comp],minlength=len(sim.kc)).astype(np.float32)
    dan=sim.dan_targets(cfg['matched_DAN'],cfg['target_compartment']);duration=cfg['duration_ms'];nsteps=round(duration/sim.dt)
    identity.to_parquet(exp/'protocol/frozen-edge-identity.parquet',index=False)
    calrows=[];cached={};wall0=time.perf_counter();cpu0=resource.getrusage(resource.RUSAGE_SELF);torch.cuda.reset_peak_memory_stats()
    for oi,(_,odor) in enumerate(cal.iterrows()):
        label='CAL_'+chr(65+oi);targets,rates=make_drive(ocfg,panel,vectors,selection,mapping,odor.odor_id,'physiological')
        pd.DataFrame({'neuron_idx':targets,'rate_hz':rates}).to_parquet(exp/'protocol'/f'{label}-drive.parquet',index=False)
        for seed in cfg['calibration_seeds']:
            actual=seed+oi*cfg['odor_seed_offset'];r=sim.run(targets,rates,actual,duration,collect_events=True);assert r['finite']
            cached[(label,seed)]=(targets,rates,actual,r);persist_run(exp,f'cal-{label}-{seed}',r)
            cs=curves(r,lookup,multi,cfg['windows_ms'],sim.dt,nsteps)
            for i,w in enumerate(cfg['windows_ms']):
                den=float(cs[i].mean())
                for u in cfg['candidate_US_ms']:
                    m=float(cs[i][round(u/sim.dt)])
                    calrows.append({'label':label,'seed':seed,'window_ms':w,'US_ms':u,'matched':m,'random_expectation':den,'ratio':m/den if den else None})
    pd.DataFrame(calrows).to_parquet(exp/'calibration/grid.parquet',index=False)
    grid=[timing_test(calrows,cfg,w,u) for w in cfg['windows_ms'] for u in cfg['candidate_US_ms']]
    good=[x for x in grid if x['pass']];chosen=min(good,key=lambda x:(x['window_ms'],x['US_ms'])) if good else None
    # Persist the selection BEFORE accessing heldout data or formal odor responses.
    save(exp/'calibration/selection.json',{'grid':grid,'chosen':chosen,'pass':chosen is not None,'formal_AB_used':False})
    summary={'phase':'4S','calibration':chosen,'heldout':None,'formal':None,'isolated':None,'whole':None,
             'timing_candidate':False,'credit_candidate':False,'associative_memory_claim':False}
    if chosen:
        w=chosen['window_ms'];u=chosen['US_ms'];us=round(u/sim.dt);tau=w/3
        held=[];gate_exact=[]
        for oi,(_,odor) in enumerate(cal.iterrows()):
            label='CAL_'+chr(65+oi);targets,rates=make_drive(ocfg,panel,vectors,selection,mapping,odor.odor_id,'physiological')
            for seed in cfg['heldout_seeds']:
                actual=seed+oi*cfg['odor_seed_offset'];r=sim.run(targets,rates,actual,duration,collect_events=True);assert r['finite']
                cs=curves(r,lookup,multi,[w],sim.dt,nsteps)[0];den=float(cs.mean());m=float(cs[us])
                held.append({'label':label,'seed':seed,'window_ms':w,'US_ms':u,'matched':m,'random_expectation':den,'ratio':m/den if den else None})
                persist_run(exp,f'held-{label}-{seed}',r)
                for pulse in (us,int(np.random.default_rng(actual+975000).integers(nsteps))):
                    forced=sim.run(targets,rates,actual,duration,dan_targets=dan,dan_pulses=[pulse*sim.dt],collect_events=True)
                    t,_=last_spike_trace(r['kc_event_step'],r['kc_event_node'],lookup,pulse,w,tau,sim.dt)
                    tt,_=last_spike_trace(forced['kc_event_step'],forced['kc_event_node'],lookup,pulse,w,tau,sim.dt)
                    v2check=v2_trace(forced['kc_event_step'],forced['kc_event_node'],lookup,pulse,sim.dt)
                    gate_exact.append({'label':label,'seed':seed,'US_step':pulse,'credit_exact':bool(np.array_equal(t,tt)),
                                       'v2_CPU_trace_exact_to_runtime':bool(np.array_equal(v2check,forced['trace']))})
        pd.DataFrame(held).to_parquet(exp/'heldout/results.parquet',index=False)
        h=timing_test(held,cfg,w,u);h['same_event_vs_live_US_credit_exact']=all(x['credit_exact'] and x['v2_CPU_trace_exact_to_runtime'] for x in gate_exact)
        h['pass']=h['pass'] and h['same_event_vs_live_US_credit_exact'];save(exp/'heldout/gate.json',h);save(exp/'regression/live-US-gates.json',gate_exact)
        summary['heldout']=h;print(json.dumps({'selection':chosen,'heldout':h}),flush=True)
        if h['pass']:
            summary['timing_candidate']=True
            formal(cfg,exp,sim,identity,lookup,loc,comp,w,u,ocfg,panel,vectors,selection,mapping,summary)
    summary['regression_pass']=evidence_regression(cfg,exp,rt,sim)
    summary['status']='TIMING-PASS-CREDIT-PASS' if summary['credit_candidate'] else ('TIMING-PASS-CREDIT-NOT-ESTABLISHED' if summary['timing_candidate'] else 'FAIL-TIMING-VALIDATION')
    summary['timing_candidate'] &= summary['regression_pass'];summary['credit_candidate'] &= summary['regression_pass']
    save(exp/'audit/results.json',summary)
    cpu1=resource.getrusage(resource.RUSAGE_SELF);wall=time.perf_counter()-wall0
    save(exp/'performance/summary.json',{'wall_seconds':wall,'peak_VRAM_bytes':int(torch.cuda.max_memory_allocated()),
      'peak_RAM_bytes':int(cpu1.ru_maxrss*1024),'CPU_percent_one_core':100*(cpu1.ru_utime+cpu1.ru_stime-cpu0.ru_utime-cpu0.ru_stime)/wall,
      'GPU':'RTX 4090','GPU_utilization':'not sampled','runtime':'deterministic-runtime-v2'})
    print(json.dumps(summary),flush=True)


def formal(cfg,exp,sim,identity,lookup,loc,comp,w,u,ocfg,panel,vectors,selection,mapping,summary):
    spec=cfg['formal'];us=round(u/sim.dt);nsteps=round(cfg['duration_ms']/sim.dt);names=spec['arms'];points=spec['checkpoints']
    targets,rates=make_drive(ocfg,panel,vectors,selection,mapping,selection['ODOR_A']['odor_id'],'physiological')
    states={};rows=[];audit=[];contam=[];repeat=[];groups=identity.identity_group.astype(str).to_numpy()
    for seed in spec['train_seed_bases']:
        E={name:np.ones(len(identity),np.float32) for name in names}
        for name in names:
            states[f'{seed}_{name}_0']=E[name].copy()
            for x in group_metrics(E[name],sim.base_candidate,identity):rows.append({'seed':seed,'pairing':0,'arm':name,**x})
        for pairing in range(1,max(points)+1):
            actual=seed+pairing;r=sim.run(targets,rates,actual,cfg['duration_ms'],collect_events=True);assert r['finite']
            s=r['kc_event_step'];n=r['kc_event_node'];rand=int(np.random.default_rng(actual+spec['random_seed_offset']).integers(nsteps))
            e3,last=last_spike_trace(s,n,lookup,us,w,w/3,sim.dt);er,lastr=last_spike_trace(s,n,lookup,rand,w,w/3,sim.dt)
            v2=v2_trace(s,n,lookup,us,sim.dt);v90=v2_trace(s,n,lookup,900,sim.dt)
            elig={'v2_same_US':edge_elig(v2,loc,comp),'v2_90ms':edge_elig(v90,loc,comp),
              'v3_matched':edge_elig(e3,loc,comp),'v3_random':edge_elig(er,loc,comp),'D0':np.zeros(len(identity),np.float32)}
            event_hash=persist_run(exp,f'formal-{seed}-{pairing}',r)
            audit.append({'seed':seed,'pairing':pairing,'actual_seed':actual,'random_US_step':rand,'matched_US_step':us,'KC_hash_all_arms':event_hash,
                          'no_weight_overlay_during_event_collection':True,'same_KC_events_all_arms':True})
            for name in names:
                E[name],delta=update_E(E[name],elig[name])
                for j in np.flatnonzero((groups=='B-specific')&(delta!=0)):
                    gate=rand if name=='v3_random' else (900 if name=='v2_90ms' else us)
                    ts=s[n==int(identity.source_neuron_idx.iloc[j])];ts=ts[ts<=gate];ls=int(ts[-1]) if len(ts) else None
                    contam.append({'seed':seed,'pairing':pairing,'arm':name,'edge_id':int(identity.edge_id.iloc[j]),
                      'KC_body_id':int(identity.source_body_id.iloc[j]),'last_KC_step':ls,'US_step':gate,'eligibility':float(elig[name][j]),'delta_E':float(delta[j])})
            if pairing in points:
                for name in names:
                    states[f'{seed}_{name}_{pairing}']=E[name].copy()
                    for x in group_metrics(E[name],sim.base_candidate,identity):rows.append({'seed':seed,'pairing':pairing,'arm':name,**x})
                # Checkpoint each independent trajectory for recoverability.
                np.savez_compressed(exp/'formal'/f'states-{seed}.npz',**{k:v for k,v in states.items() if k.startswith(str(seed)+'_')})
            if pairing%20==0:print(json.dumps({'formal_seed':seed,'pairings':pairing}),flush=True)
        if seed==spec['train_seed_bases'][0]:
            rr=sim.run(targets,rates,seed+100,cfg['duration_ms'],collect_events=True)
            repeat.append({'seed':seed+100,'events_exact':arrhash(rr['kc_event_step'],rr['kc_event_node'])==event_hash,
                           'states_exact':arrhash(rr['final_v'],rr['final_g'])==arrhash(r['final_v'],r['final_g'])})
    f=pd.DataFrame(rows);f.to_parquet(exp/'formal/group-trajectory.parquet',index=False)
    pd.DataFrame(audit).to_parquet(exp/'formal/event-audit.parquet',index=False);pd.DataFrame(contam).to_parquet(exp/'formal/B-contamination.parquet',index=False)
    save(exp/'regression/formal-repeat.json',repeat)
    endpoints=f[f.pairing==spec['primary_endpoint']];gates=[]
    for seed in spec['train_seed_bases']:
        z=endpoints[endpoints.seed==seed]
        def val(arm,group):return float(z[(z.arm==arm)&(z.group==group)].total_abs_delta_E.iloc[0])
        ma=val('v3_matched','A-specific');ra=val('v3_random','A-specific');b2=val('v2_same_US','B-specific');b3=val('v3_matched','B-specific')
        ratio=ma/ra if ra else None;red=(1-b3/b2) if b2 else None
        ag=ma>0 and (ra==0 or ma>=spec['minimum_A_credit_matched_random_ratio']*ra)
        bg=(red is not None and red>=spec['minimum_B_reduction_vs_v2_same_US']) or (b2==0 and b3==0)
        gates.append({'seed':seed,'A_credit':ma,'random_A_credit':ra,'A_ratio':ratio,'B_v2':b2,'B_v3':b3,'B_reduction':red,'pass':ag and bg})
    fp=sum(x['pass'] for x in gates)>=spec['minimum_seed_majority']
    summary['formal']={'gate_pass':fp,'primary_endpoint':spec['primary_endpoint'],'seeds':gates,'open_loop_shadow_only':True}
    save(exp/'formal/gate.json',summary['formal'])
    # Always measure isolated transfer after independently validated timing; its own gate controls whole-brain overlays.
    replay=IsolatedReplay(sim,cfg['target_compartment']);evals={};results=[]
    for odor in ('ODOR_A','ODOR_B'):
        t,rates=make_drive(ocfg,panel,vectors,selection,mapping,selection[odor]['odor_id'],'physiological')
        for es in spec['evaluation_seeds']:
            er=sim.run(t,rates,es,cfg['duration_ms'],collect_events=True);evals[(odor,es)]=er;persist_run(exp,f'eval-{odor}-{es}',er)
    for pairing in spec['isolated_checkpoints']:
        for odor in ('ODOR_A','ODOR_B'):
            for seed,es in zip(spec['train_seed_bases'],spec['evaluation_seeds']):
                er=evals[(odor,es)];base=replay.run(er['kc_event_step'],er['kc_event_node'],None,cfg['duration_ms'])
                for name in names:
                    post=replay.run(er['kc_event_step'],er['kc_event_node'],states[f'{seed}_{name}_{pairing}'],cfg['duration_ms'])
                    results.append({'pairing':pairing,'odor':odor,'seed':seed,'evaluation_seed':es,'arm':name,**replay_metrics(base,post,sim.dt)})
    z=pd.DataFrame(results);z.to_parquet(exp/'isolated/results.parquet',index=False);ig=[];trajectory=[]
    for pairing in spec['isolated_checkpoints']:
        for seed in spec['train_seed_bases']:
            def L(arm,odor):return int(z[(z.pairing==pairing)&(z.seed==seed)&(z.arm==arm)&(z.odor==odor)].spike_L1.iloc[0])
            a,b,ra,rb=L('v3_matched','ODOR_A'),L('v3_matched','ODOR_B'),L('v3_random','ODOR_A'),L('v3_random','ODOR_B')
            trajectory.append({'pairing':pairing,'seed':seed,'A':a,'B':b,'random_A':ra,'random_B':rb,'A_specific':a>0 and a>b,'random_specific':ra>0 and ra>rb,'matched_beats_random':a>ra})
    endpoint=[x for x in trajectory if x['pairing']==spec['primary_endpoint']];major=spec['minimum_seed_majority']
    ip=sum(x['A_specific'] for x in endpoint)>=major and sum(x['random_specific'] for x in endpoint)<major and sum(x['matched_beats_random'] for x in endpoint)>=major and bool((z[z.arm=='D0'].spike_L1==0).all())
    summary['isolated']={'gate_pass':ip,'endpoint':endpoint,'trajectory':trajectory,'D0_zero':bool((z[z.arm=='D0'].spike_L1==0).all()),'no_checkpoint_selected_posthoc':True}
    save(exp/'isolated/gate.json',summary['isolated'])
    if fp and ip:
        whole=[]
        for odor in ('ODOR_A','ODOR_B'):
            t,rates=make_drive(ocfg,panel,vectors,selection,mapping,selection[odor]['odor_id'],'physiological')
            for seed,es in zip(spec['train_seed_bases'],spec['evaluation_seeds']):
                base=evals[(odor,es)]
                for name in ('v3_matched','v3_random','v2_same_US','D0'):
                    post=sim.run(t,rates,es,cfg['duration_ms'],E=states[f'{seed}_{name}_100'])
                    whole.append({'seed':seed,'odor':odor,'arm':name,'MBON_L1':int(np.abs(post['mbon_counts']-base['mbon_counts']).sum()),
                      'downstream_L1':int(np.abs(post['downstream_counts']-base['downstream_counts']).sum()),'finite':post['finite']})
        wf=pd.DataFrame(whole);wf.to_parquet(exp/'whole-brain/results.parquet',index=False);wg=[]
        for seed in spec['train_seed_bases']:
            def v(arm,odor):return int(wf[(wf.seed==seed)&(wf.arm==arm)&(wf.odor==odor)].MBON_L1.iloc[0])
            a,b,ra,rb=v('v3_matched','ODOR_A'),v('v3_matched','ODOR_B'),v('v3_random','ODOR_A'),v('v3_random','ODOR_B')
            wg.append({'seed':seed,'A':a,'B':b,'random_A':ra,'random_B':rb,'specific':a>b and a>ra,'random_specific':ra>rb})
        wp=sum(x['specific'] for x in wg)>=major and sum(x['random_specific'] for x in wg)<major and bool(wf.finite.all()) and bool((wf[wf.arm=='D0'].MBON_L1==0).all())
        summary['whole']={'run':True,'pass':wp,'seeds':wg};summary['credit_candidate']=bool(wp)
    else:summary['whole']={'run':False,'reason':'formal credit or isolated specificity gate failed'}
    save(exp/'whole-brain/gate.json',summary['whole'])


def finalize(cfg,exp,model):
    old=phase4p_finalize.integrity();now=checks(cfg);before=json.loads((exp/'audit/preflight.json').read_text())
    pre=json.loads((exp/'protocol/preregistered.json').read_text());same=all(digest(model/n)==h for n,h in pre['source_hashes'].items())
    integrity={'Phase1_4O':old,'Phase4P_4R':now,'before_after_equal':now==before,'preregistered_source_hashes_unchanged':same,
               'pass':old['all_static_checks_passed'] and all(x['passed'] for x in now.values()) and now==before and same}
    save(exp/'audit/frozen-integrity.json',integrity);summary=json.loads((exp/'audit/results.json').read_text())
    summary['integrity_pass']=integrity['pass'];summary['timing_candidate'] &= integrity['pass'];summary['credit_candidate'] &= integrity['pass']
    if summary['timing_candidate']:
        save(model/'MaleCNS-CSUS-Timing-v1-candidate.json',{'scope':'independent calibration and seed-heldout matched/random credit timing only; no memory claim',
             'selection':summary['calibration'],'heldout':summary['heldout'],'brain_unchanged':True})
    if summary['credit_candidate']:
        save(model/'MaleCNS-CreditAssignment-v1-candidate.json',{'scope':'open-loop shadow credit and replay; no online learning claim','source':str(exp/'audit/results.json')})
    save(exp/'audit/results.json',summary)
    for root in (model,exp):
        ps=[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 ps]
        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))
        test=verify_checksum(root,root/'checksums.sha256');assert test['passed'];print(json.dumps(test),flush=True)
    print(json.dumps({'status':summary['status'],'timing_candidate':summary['timing_candidate'],'credit_candidate':summary['credit_candidate'],'integrity_pass':integrity['pass']}),flush=True)


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