"""Preregistered independent fixed-US stability test. Frozen brain, no learning."""
from pathlib import Path
import argparse
import json
import time
import numpy as np
import pandas as pd
import torch
import yaml
import phase4p
import phase4p_finalize
import phase4s
from phase4p_revalidate import DeterministicAPLSim,door_inputs,make_drive
from phase4r import save,digest,arrhash,curves
from phase4q import last_spike_trace,v2_trace,edge_elig
from phase4h_audit import verify_checksum

def checks(cfg):
    sc=yaml.safe_load((Path(cfg['source_model'])/'phase4s.yaml').read_text())
    out=phase4s.checks(sc)
    for key in ('source_root','source_model'):
        root=Path(cfg[key]);out[key]=verify_checksum(root,root/'checksums.sha256')
    assert all(x['passed'] for x in out.values())
    return out

def gate(rows,cfg,w,u):
    f=pd.DataFrame(rows);f=f[(f.window_ms==w)&(f.US_ms==u)];out=[]
    for odor,g in f.groupby('odor_id',sort=True):
        den=float(g.random_expectation.mean());rat=float(g.matched.mean()/den) if den else None
        q=float(np.quantile(g.ratio,0.2));n=int((g.ratio>=cfg['timing_gate']['ratio']).sum())
        out.append({'odor_id':odor,'n':len(g),'qualifying':n,'ratio_of_means':rat,'q20_ratio':q,
          'zero_credit':int((g.matched==0).sum()),'mean_credit':float(g.matched.mean()),
          'pass':bool(den>0 and rat>=cfg['timing_gate']['mean_ratio_each_odor'] and n/len(g)>=cfg['timing_gate']['qualifying_fraction_each_odor'])})
    return {'window_ms':int(w),'US_ms':int(u),'tau_ms':float(w/3),'odors':out,
        'score':min(x['q20_ratio'] for x in out),'pass':bool(out and all(x['pass'] for x in out))}

def controls(w,u,dt,root):
    lookup=np.arange(3,dtype=np.int32);us=round(u/dt);tau=w/3
    def calc(ss,nn,gate=us):return last_spike_trace(np.array(ss,np.int32),np.array(nn,np.int32),lookup,gate,w,tau,dt)[0]
    recent=calc([us-1],[0]);future=calc([us+1],[0]);old=calc([0],[0],round((w+10)/dt))
    unrelated=edge_elig(recent,np.array([0,1,2]),np.zeros(3,bool))
    d0=edge_elig(recent,np.array([0,1,2]),np.ones(3,bool),teaching=False)
    base=calc([us-1,us],[0,1]);shift=last_spike_trace(np.array([us+17,us+18]),np.array([0,1]),lookup,us+18,w,tau,dt)[0]
    path=root/'trace-checkpoint.npz';np.savez(path,trace=base)
    tests={'recent_positive':bool(recent[0]>0),'inactive_zero':bool(recent[1]==0),
       'future_zero':bool(not future.any()),'distant_zero':bool(not old.any()),
       'unrelated_zero':bool(not unrelated.any()),'D0_zero':bool(not d0.any()),
       'same_step_one':bool(base[1]==1),'both_arrivals_shift_exact':bool(np.array_equal(base,shift)),
       'disk_checkpoint_exact':bool(np.array_equal(base,np.load(path)['trace'])),
       'repeated_algebra_exact':all(np.array_equal(base,calc([us-1,us],[0,1])) for _ in range(10))}
    save(root/'unit-tests.json',{'scope':'Algebraic controls; not physiological simulations','tests':tests,'pass':all(tests.values())})
    assert all(tests.values())

def run(cfg,exp,model):
    for d in ('protocol','calibration','validation','events','regression','performance'):(exp/d).mkdir(parents=True,exist_ok=True)
    assert not (exp/'protocol/preregistered.json').exists(),'Refusing overwrite'
    audit=json.loads((exp/'audit/summary.json').read_text());assert audit['all_zero_trials_explained_by_event_timing']
    pre=checks(cfg);save(exp/'protocol/preflight.json',pre)
    ocfg,panel,vectors,selection,mapping=door_inputs()
    cal=panel[panel.split=='calibration'].iloc[:2]
    formal={selection['ODOR_A']['odor_id'],selection['ODOR_B']['odor_id']}
    novel=panel[(panel.split=='heldout')&~panel.odor_id.isin(formal)]
    assert len(cal)==2 and len(novel)==4 and not set(cal.odor_id)&set(novel.odor_id)
    odors=pd.concat([cal,novel]).reset_index(drop=True)
    seedsets=[{int(s)+i*cfg['odor_seed_offset'] for i in range(len(odors)) for s in cfg[k]} for k in ('calibration_seeds','validation_seeds')]
    assert seedsets[0].isdisjoint(seedsets[1])
    prereg={'config':cfg,'odors':odors[['odor_id','odor_name','split']].to_dict('records'),
      'created_before_new_simulations':True,'sources':{p.name:digest(p) for p in model.glob('phase4t*') if p.is_file()},
      'historical_4S_used_only_for_hypothesis':True,'new_heldout_not_used_for_selection':True}
    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)
    lookup=np.full(sim.n,-1,np.int32);lookup[sim.kc]=np.arange(len(sim.kc))
    identity=sim.identity;assert len(identity)==14232
    target=identity.compartment.eq(cfg['target_compartment']).to_numpy()
    loc=lookup[identity.source_neuron_idx.to_numpy()]
    multi=np.bincount(loc[target],minlength=len(sim.kc)).astype(np.float32)
    duration=cfg['duration_ms'];nsteps=round(duration/sim.dt);dan=sim.dan_targets(cfg['matched_DAN'],cfg['target_compartment'])
    start=time.perf_counter();torch.cuda.reset_peak_memory_stats();records=[];live=[];repeats=[]
    def trial(oi,seed,split,ws,us):
        odor=odors.iloc[oi];t,rates=make_drive(ocfg,panel,vectors,selection,mapping,odor.odor_id,'physiological')
        actual=seed+oi*cfg['odor_seed_offset']
        r=sim.run(t,rates,actual,duration,collect_events=True);assert r['finite']
        phase4s.persist_run(exp,f'{split}-{oi}-{seed}',r)
        c=curves(r,lookup,multi,ws,sim.dt,nsteps)
        for wi,w in enumerate(ws):
            den=float(c[wi].mean())
            for u in us:
                m=float(c[wi][round(u/sim.dt)])
                records.append({'split':split,'odor_id':odor.odor_id,'odor_index':oi,'seed':seed,'actual_seed':actual,
                    'window_ms':w,'US_ms':u,'matched':m,'random_expectation':den,'ratio':m/den if den else 0.})
        return r,t,rates,actual
    for oi in range(len(cal)):
        for seed in cfg['calibration_seeds']:trial(oi,seed,'calibration',cfg['windows_ms'],cfg['candidate_US_ms'])
        print(json.dumps({'stage':'calibration','odor_index':oi,'seeds':len(cfg['calibration_seeds'])}),flush=True)
    grid=[gate(records,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=sorted(good,key=lambda x:(-x['score'],x['window_ms'],x['US_ms']))[0] if good else None
    pd.DataFrame(records).to_parquet(exp/'calibration/grid.parquet',index=False)
    save(exp/'calibration/selection.json',{'grid':grid,'chosen':chosen,'selected_before_validation':True})
    print(json.dumps({'selected':chosen}),flush=True)
    summary={'phase':'4T','P0_diagnosis_pass':True,'calibration':chosen,'seed_validation':None,'odor_validation':None,
      'timing_gate_pass':False,'credit_specificity_pass':False,'formal_training_run':False,'memory_claim':False}
    if chosen:
        w=chosen['window_ms'];u=chosen['US_ms'];us=round(u/sim.dt);controls(w,u,sim.dt,exp/'regression')
        for oi in range(len(odors)):
            split='seed_validation' if oi<len(cal) else 'odor_validation'
            for si,seed in enumerate(cfg['validation_seeds']):
                r,t,rates,actual=trial(oi,seed,split,[w],[u])
                if si<cfg['live_US_checks']['first_seeds_per_validation_odor']:
                    for pulse in (us,int(np.random.default_rng(actual+cfg['live_US_checks']['random_seed_offset']).integers(nsteps))):
                        forced=sim.run(t,rates,actual,duration,dan_targets=dan,dan_pulses=[pulse*sim.dt],collect_events=True);assert forced['finite']
                        phase4s.persist_run(exp,f'live-{oi}-{seed}-{pulse}',forced)
                        a,_=last_spike_trace(r['kc_event_step'],r['kc_event_node'],lookup,pulse,w,w/3,sim.dt)
                        b,_=last_spike_trace(forced['kc_event_step'],forced['kc_event_node'],lookup,pulse,w,w/3,sim.dt)
                        v=v2_trace(forced['kc_event_step'],forced['kc_event_node'],lookup,pulse,sim.dt)
                        live.append({'odor_index':oi,'seed':seed,'US_step':pulse,'trace_exact':bool(np.array_equal(a,b)),
                            'v2_trace_exact':bool(np.array_equal(v,forced['trace']))})
                if si==0:
                    rr=sim.run(t,rates,actual,duration,collect_events=True);assert rr['finite']
                    repeats.append({'odor_index':oi,'seed':seed,'events_exact':arrhash(rr['kc_event_step'],rr['kc_event_node'])==arrhash(r['kc_event_step'],r['kc_event_node']),
                        'states_exact':arrhash(rr['final_v'],rr['final_g'])==arrhash(r['final_v'],r['final_g'])})
            print(json.dumps({'stage':split,'odor_index':oi,'seeds':len(cfg['validation_seeds'])}),flush=True)
        rows=[r for r in records if r['split']!='calibration'];pd.DataFrame(rows).to_parquet(exp/'validation/results.parquet',index=False)
        for split in ('seed_validation','odor_validation'):summary[split]=gate([r for r in rows if r['split']==split],cfg,w,u)
        save(exp/'regression/live-US-checks.json',live);save(exp/'regression/repeat-checks.json',repeats)
        summary['live_controls_pass']=all(x['trace_exact'] and x['v2_trace_exact'] for x in live)
        summary['repeats_pass']=all(x['events_exact'] and x['states_exact'] for x in repeats)
        summary['timing_gate_pass']=summary['seed_validation']['pass'] and summary['odor_validation']['pass'] and summary['live_controls_pass'] and summary['repeats_pass']
    summary['frozen_reference_pass']=phase4s.evidence_regression(cfg,exp,rt,sim)
    summary['timing_gate_pass'] &= summary['frozen_reference_pass']
    summary['status']='PASS-INDEPENDENT-TIMING-ONLY' if summary['timing_gate_pass'] else 'FAIL-TIMING-STABILITY-OR-GENERALIZATION'
    summary['next_gate']='Independent A/B local-credit specificity; no memory claim' if summary['timing_gate_pass'] else 'Stop downstream learning; report failed timing/odor gates without retuning'
    save(exp/'validation/summary.json',summary)
    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),'GPU_utilization':'not sampled','weight_updates':0})
    print(json.dumps(summary,indent=2),flush=True)

def seal(cfg,exp,model):
    before=json.loads((exp/'protocol/preflight.json').read_text());after=checks(cfg)
    old=phase4p_finalize.integrity();p=json.loads((exp/'protocol/preregistered.json').read_text())
    same=all(digest(model/n)==h for n,h in p['sources'].items())
    integrity={'old':old,'current':after,'before_after_equal':before==after,'source_hashes_unchanged':same,
       'pass':old['all_static_checks_passed'] and before==after and same}
    save(exp/'audit/frozen-integrity.json',integrity);assert integrity['pass']
    summary=json.loads((exp/'validation/summary.json').read_text());summary['integrity_pass']=True
    if summary['timing_gate_pass']:
        save(model/'MaleCNS-CSUS-Timing-Stability-v1-candidate.json',{'scope':'Independent seed and timing-heldout odor validation only; not local credit specificity or memory',
           'selection':summary['calibration'],'validation':summary})
    save(exp/'validation/summary.json',summary)
    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 ('checksums.sha256','manifest.json')]
        save(root/'manifest.json',{'files':rows});rows.append({'path':'manifest.json','sha256':digest(root/'manifest.json')})
        (root/'checksums.sha256').write_text(''.join(r['sha256']+'  '+r['path']+'\n' for r in rows))
        result=verify_checksum(root,root/'checksums.sha256');assert result['passed'];print(json.dumps(result),flush=True)

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