"""Literature-informed panel coverage test; frozen brain and encoder, 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
import phase4t
from phase4p_revalidate import DeterministicAPLSim,door_inputs,make_drive
from phase4l import load_external
from phase4r import save,digest,arrhash,curves
from phase4q import last_spike_trace,v2_trace
from phase4h_audit import verify_checksum

def checks(cfg):
    tc=yaml.safe_load((Path(cfg['prior_model'])/'phase4t.yaml').read_text())
    out=phase4t.checks(tc)
    for key in ('prior_root','prior_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 prepare(cfg,exp):
    oc,old,oldvec,selection,mapping=door_inputs()
    _,response,odors,provenance=load_external(oc,exp)
    assert 'InChIKey' in odors.columns and 'Name' in odors.columns
    usable=mapping[mapping.mapping_confidence.isin(['exact','alias-normalized'])&mapping.neuron_idx.notna()]
    cg=usable[['odor_channel','glomerulus']].drop_duplicates();gloms=selection['glomeruli_order'];gi={g:i for i,g in enumerate(gloms)}
    assert sorted(cg.glomerulus.unique())==gloms
    names=odors.dropna(subset=['InChIKey']).drop_duplicates('InChIKey').set_index('InChIKey')['Name'].to_dict()
    vectors={};rows=[]
    for key,row in response.iterrows():
        if str(key)=='SFR':continue
        name=names.get(key)
        if pd.isna(name) or not str(name).strip():continue
        arr=np.zeros(len(gloms),np.float32)
        for _,q in cg.iterrows():
            value=row.get(q.odor_channel)
            if pd.notna(value):arr[gi[q.glomerulus]]=max(arr[gi[q.glomerulus]],max(float(value),0.))
        allsum=float(np.maximum(np.nan_to_num(row.to_numpy(np.float64),nan=0.),0.).sum())
        mapped=float(arr.sum());strong=int(np.count_nonzero(arr>=float(oc['odor_encoder']['strong_response_threshold'])))
        if mapped<=0 or strong<int(oc['odor_encoder']['minimum_mapped_strong_channels']):continue
        oid=str(key);vectors[oid]=arr
        rows.append({'odor_id':oid,'odor_name':str(name),'mapped_response_sum':mapped,'all_response_sum':allsum,
            'mapped_signal_fraction':mapped/allsum if allsum else 0.,'mapped_strong_glomeruli':strong,'mapped_nonzero_glomeruli':int(np.count_nonzero(arr))})
    candidates=pd.DataFrame(rows)
    tests=[]
    for oid in old.odor_id:
        assert np.array_equal(vectors[oid],oldvec[oid]),'Frozen vector mismatch: '+oid
        # make_drive is the unchanged encoder; physiological mode never rescales to the panel median.
        t,r=make_drive(oc,old,vectors,selection,mapping,oid,'physiological')
        tt,rr=make_drive(oc,old,oldvec,selection,mapping,oid,'physiological')
        assert np.array_equal(t,tt) and np.array_equal(r,rr)
        tests.append({'odor_id':oid,'vector_exact':True,'PN_targets_rates_exact':True,'rate_sum_hz':float(r.sum())})
    oldnames=set(old.odor_name.astype(str).str.strip().str.casefold())
    pool=candidates[~candidates.odor_id.isin(old.odor_id)&~candidates.odor_name.str.strip().str.casefold().isin(oldnames)&(candidates.mapped_signal_fraction>=.9)].copy()
    pool=pool.sort_values(['mapped_response_sum','odor_id'],kind='stable').drop_duplicates('odor_name').reset_index(drop=True)
    n=int(cfg['panel']['novel_count']);assert len(pool)>=n,'Not enough independent physiology odors'
    bins=np.array_split(np.arange(len(pool)),n);selected=[int(b[(len(b)-1)//2]) for b in bins]
    novel=pool.iloc[selected].copy();novel['split']='new_odor_validation'
    cal=old[old.split=='calibration'].drop(columns=['formal_label']).copy();cal['split']='calibration';assert len(cal)==6
    panel=pd.concat([cal,novel],ignore_index=True)
    formal={selection['ODOR_A']['odor_id'],selection['ODOR_B']['odor_id']}
    assert not set(panel.odor_id)&formal and not set(novel.odor_id)&set(old.odor_id)
    pool.to_parquet(exp/'protocol/eligible-new-odor-pool.parquet',index=False)
    panel.to_parquet(exp/'protocol/frozen-panel.parquet',index=False)
    np.savez_compressed(exp/'protocol/frozen-vectors.npz',glomeruli=np.asarray(gloms),**{oid:vectors[oid] for oid in panel.odor_id})
    save(exp/'regression/frozen-encoder-exact.json',{'pass':True,'rows':tests})
    save(exp/'protocol/panel-selection.json',{'eligible_new_odors':len(pool),'rank_bins':[list(map(int,b)) for b in bins],
      'selected_ranks':selected,'new_odor_ID_overlap_old12':0,'input_response_only':True,'runtime_responses_used':False,
      'source':provenance,'frozen_mapping_sha256':digest(Path(oc['paths']['phase4l_root'])/'receptor-mapping/olfactory-mapping.parquet')})
    return oc,old,panel,vectors,selection,mapping

def run(cfg,exp,model):
    for d in ('protocol','external-resources','calibration','validation','events','regression','audit','performance'):(exp/d).mkdir(parents=True,exist_ok=True)
    assert not (exp/'protocol/preregistered.json').exists(),'Refusing overwrite'
    before=checks(cfg);save(exp/'audit/preflight.json',before)
    oc,old,panel,vectors,selection,mapping=prepare(cfg,exp)
    cal=panel[panel.split=='calibration'];nc=len(cal)
    save(exp/'protocol/preregistered.json',{'config':cfg,'panel':panel.to_dict('records'),
      'sources':{p.name:digest(p) for p in model.glob('phase4u*') if p.is_file()},
      'created_before_simulation':True,'old12_excluded_from_new_test':True,'no_new_dynamic_results_seen':True,
      'method':'Expanded calibration coverage; fixed Phase4T grid, score and gate. Not posthoc retest of Phase4T failures.'})
    print(json.dumps({'panel':panel[['odor_id','odor_name','split','mapped_response_sum']].to_dict('records')}),flush=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)
    lookup=np.full(sim.n,-1,np.int32);lookup[sim.kc]=np.arange(len(sim.kc))
    identity=sim.identity;assert len(identity)==14232
    mask=identity.compartment.eq(cfg['target_compartment']).to_numpy();loc=lookup[identity.source_neuron_idx.to_numpy(np.int32)]
    multi=np.bincount(loc[mask],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)
    records=[];repeats=[];live=[];times=[];start=time.perf_counter();torch.cuda.reset_peak_memory_stats()
    def trial(oi,seed,split,ws,us):
        odor=panel.iloc[oi];targets,rates=make_drive(oc,old,vectors,selection,mapping,odor.odor_id,'physiological')
        actual=cfg['seed_namespace_offset']+oi*cfg['odor_seed_offset']+seed
        r=sim.run(targets,rates,actual,duration,collect_events=True);assert r['finite']
        phase4s.persist_run(exp,f'{split}-{oi}-{seed}',r)
        cs=curves(r,lookup,multi,ws,sim.dt,nsteps)
        for wi,w in enumerate(ws):
            den=float(cs[wi].mean())
            for u in us:
                m=float(cs[wi][round(u/sim.dt)])
                records.append({'split':split,'odor_id':odor.odor_id,'odor_name':odor.odor_name,'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.})
        s=r['kc_event_step'];nodes=r['kc_event_node'];target=multi[lookup[nodes]]>0;ts=s[target]
        times.append({'split':split,'odor_index':oi,'seed':seed,'KC_active':int(np.count_nonzero(r['kc_counts'])),
           'target_active_KC':int(np.unique(nodes[target]).size),'target_events':len(ts),
           'first_target_ms':float(ts.min()*sim.dt) if len(ts) else None,
           'median_target_ms':float(np.median(ts)*sim.dt) if len(ts) else None})
        return r,targets,rates,actual
    for oi in range(nc):
        for si,seed in enumerate(cfg['calibration_seeds']):
            r,t,rates,actual=trial(oi,seed,'calibration',cfg['windows_ms'],cfg['candidate_US_ms'])
            if si==0:
                rr=sim.run(t,rates,actual,duration,collect_events=True);assert rr['finite']
                repeats.append({'split':'calibration','odor_index':oi,'seed':seed,
                   'events_exact':arrhash(r['kc_event_step'],r['kc_event_node'])==arrhash(rr['kc_event_step'],rr['kc_event_node']),
                   'states_exact':arrhash(r['final_v'],r['final_g'])==arrhash(rr['final_v'],rr['final_g'])})
        print(json.dumps({'stage':'calibration','odor_index':oi,'seeds':len(cfg['calibration_seeds'])}),flush=True)
    grid=[phase4t.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,'saved_before_validation':True})
    print(json.dumps({'selected':chosen}),flush=True)
    summary={'phase':'4U','calibration':chosen,'seed_validation':None,'new_odor_validation':None,'timing_gate_pass':False,
        'weight_updates':0,'memory_claim':False,'calibration_odors':nc,'new_test_odors':len(panel)-nc,'formal_AB_used':False}
    if chosen:
        w=chosen['window_ms'];u=chosen['US_ms'];us=round(u/sim.dt);phase4t.controls(w,u,sim.dt,exp/'regression')
        for oi in range(len(panel)):
            split='seed_validation' if oi<nc else 'new_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']['seeds_per_validation_odor']:
                    for pulse in (us,int(np.random.default_rng(actual+cfg['live_US_checks']['random_seed_offset']).integers(nsteps))):
                        rr=sim.run(t,rates,actual,duration,dan_targets=dan,dan_pulses=[pulse*sim.dt],collect_events=True);assert rr['finite']
                        phase4s.persist_run(exp,f'live-{oi}-{seed}-{pulse}',rr)
                        a,_=last_spike_trace(r['kc_event_step'],r['kc_event_node'],lookup,pulse,w,w/3,sim.dt)
                        b,_=last_spike_trace(rr['kc_event_step'],rr['kc_event_node'],lookup,pulse,w,w/3,sim.dt)
                        v=v2_trace(rr['kc_event_step'],rr['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,rr['trace']))})
            print(json.dumps({'stage':split,'odor_index':oi,'seeds':len(cfg['validation_seeds'])}),flush=True)
        validation=[x for x in records if x['split']!='calibration'];pd.DataFrame(validation).to_parquet(exp/'validation/results.parquet',index=False)
        for split in ('seed_validation','new_odor_validation'):summary[split]=phase4t.gate([x for x in validation if x['split']==split],cfg,w,u)
        summary['live_controls_pass']=all(x['trace_exact'] and x['v2_trace_exact'] for x in live)
        summary['timing_gate_pass']=summary['seed_validation']['pass'] and summary['new_odor_validation']['pass'] and summary['live_controls_pass']
    save(exp/'regression/live-US-checks.json',live);save(exp/'regression/repeats.json',repeats)
    summary['repeats_pass']=all(x['events_exact'] and x['states_exact'] for x in repeats)
    summary['reference_pass']=phase4s.evidence_regression(cfg,exp,rt,sim)
    summary['timing_gate_pass'] &= summary['reference_pass'] and summary['repeats_pass']
    summary['status']='PASS-INDEPENDENT-TIMING-ONLY' if summary['timing_gate_pass'] else ('FAIL-CALIBRATION-COVERAGE' if chosen is None else 'FAIL-INDEPENDENT-TIMING-GENERALIZATION')
    pd.DataFrame(times).to_parquet(exp/'audit/event-timing.parquet',index=False)
    save(exp/'audit/results.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','base_short_trials':len(times),'live_US_trials':len(live),'repeat_short_trials':len(repeats),'frozen_reference_trials':5})
    print(json.dumps(summary,indent=2),flush=True)

def seal(cfg,exp,model):
    now=checks(cfg);old=phase4p_finalize.integrity();pre=json.loads((exp/'protocol/preregistered.json').read_text())
    same=all(digest(model/n)==h for n,h in pre['sources'].items())
    integrity={'history':old,'current':now,'before_after_equal':now==json.loads((exp/'audit/preflight.json').read_text()),'source_unchanged':same}
    integrity['pass']=old['all_static_checks_passed'] and same and integrity['before_after_equal']
    save(exp/'audit/frozen-integrity.json',integrity);assert integrity['pass']
    r=json.loads((exp/'audit/results.json').read_text());r['integrity_pass']=True;save(exp/'audit/results.json',r)
    if r['timing_gate_pass']:save(model/'MaleCNS-OdorTiming-v1-candidate.json',{'scope':'engineering timing only; no credit-specificity or memory claim','results':r})
    names={x['odor_id']:x['odor_name'] for x in pre['panel']}
    lines=['# Phase 4U — 文献约束与扩展气味时序验证','',f"状态：`{r['status']}`",'',
      '文献与模型边界见模型目录 phase4u-literature.md；官方 DoOR 数据 provenance 与哈希见 external-resources/provenance.json。',
      '扩大校准集至原始六种 calibration odors；新增六种测试气味与旧十二种 ID/名称均不重叠，按照静态响应强度分层选取。',
      '十二种旧气味的 glomerular vectors、PN targets/rates 均逐位复现。固定脑、encoder 与 14,232-edge mask 未更改；无权重更新。','']
    for key in ('calibration','seed_validation','new_odor_validation'):
        if r[key]:
            g=r[key];lines +=[f"## {key}",'',f"US={g['US_ms']} ms；window={g['window_ms']} ms。",'',
                '| 气味 | 达标 seeds | mean ratio | q20 ratio | 零 credit | PASS |','|---|---:|---:|---:|---:|---|']
            for x in g['odors']:lines.append(f"| {names[x['odor_id']]} | {x['qualifying']}/{x['n']} | {x['ratio_of_means']} | {x['q20_ratio']:.4f} | {x['zero_credit']} | {x['pass']} |")
    if r['calibration'] is None:
        grid=json.loads((exp/'calibration/selection.json').read_text())['grid'];counts=[sum(x['pass'] for x in z['odors']) for z in grid]
        lines+=['没有共同达标的候选；未访问新测试气味的动态结果，也未选择最不差的参数作为正式方案。',
          f"共 {len(grid)} 个候选；单个候选最多通过 {max(counts)}/{r['calibration_odors']} 个校准气味。所有逐候选结果已保留。"]
    perf=json.loads((exp/'performance/summary.json').read_text())
    lines+=['',f"参考回归：{r['reference_pass']}；重复：{r['repeats_pass']}；历史清单完整性：{r['integrity_pass']}。",'',
      f"本轮耗时（不含编译/封存）：{perf['wall_seconds_excluding_build']:.2f}s；peak allocated VRAM：{perf['peak_VRAM_bytes']/2**20:.2f}MiB。",'',
      '没有 formal A/B credit 或 conditioning，不接游戏、不引入 RL/decoder。文献只指导覆盖、因果时序与推断边界，未提供当前模型的 MaleCNS 实测参数。',
      '结果仅适用于当前 DoOR→PN 工程输入与 LIF 模型。失败不能证明所有时序规则无效；成功也不是关联记忆成功。']
    (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))
        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);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']))