"""Raw-to-runtime provenance audit and non-mutating audit-role repair."""
from pathlib import Path
import argparse,json,time
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.feather as ft
import pyarrow.parquet as pq
import torch,yaml
import phase4p,phase4p_finalize,phase4s,phase4w
from phase4p_revalidate import DeterministicAPLSim
from phase4h_audit import verify_checksum
from phase4r import save,digest
from transmitter_roles_v1 import roles,self_test

def checks(cfg):
    old=yaml.safe_load((Path(cfg['prior_model'])/'phase4w.yaml').read_text());out=phase4w.checks(old)
    for k in ('prior_root','prior_model','prior_addendum'):
        root=Path(cfg[k]);out[k]=verify_checksum(root,root/'checksums.sha256')
    assert all(x['passed'] for x in out.values());return out

def column_equal(a,b):
    if pa.types.is_dictionary(a.type):a=a.cast(a.type.value_type)
    if pa.types.is_dictionary(b.type):b=b.cast(b.type.value_type)
    return a.combine_chunks().equals(b.combine_chunks())

def run(cfg,exp,model):
    for d in ('protocol','audit','roles','replay','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))
    root=Path(cfg['dataset_root']);files={'annotation':root/'raw/annotations/body-annotations-male-cns-v1.0-minconf-0.5.feather',
        'NT':root/'raw/annotations/body-neurotransmitters-male-cns-v1.0.feather',
        'metadata':root/'processed/neuron-only/neuron-index/unified-metadata.parquet'}
    save(exp/'protocol/preregistered.json',{'config':cfg,'sources':{p.name:digest(p) for p in model.iterdir() if p.is_file()},
        'inputs':{k:{'path':str(p),'sha256':digest(p)} for k,p in files.items()},'before_analysis':True,'no_parameter_selection':True})
    tabs={k:(pq.read_table(p) if p.suffix=='.parquet' else ft.read_table(p)) for k,p in files.items()}
    save(exp/'audit/schema.json',{k:str(v.schema) for k,v in tabs.items()})
    md=tabs['metadata'];assert md.num_rows==166700 and md['body_id'].type==pa.uint64() and md['neuron_idx'].type==pa.int32()
    assert np.array_equal(md['neuron_idx'].to_numpy(),np.arange(166700,dtype=np.int32))
    fields=[]
    for key,idfield,prefix in [('annotation','bodyId','annotation_'),('NT','body','nt_')]:
        table=tabs[key];assert table[idfield].null_count==0 and pc.count_distinct(table[idfield]).as_py()==table.num_rows
        loc=pc.index_in(md['body_id'],value_set=table[idfield].cast(pa.uint64()));aligned=table.take(loc)
        for col in table.column_names:
            exact=column_equal(aligned[col],md[prefix+col]);fields.append({'table':key,'field':col,'exact':exact})
        if key=='annotation':assert loc.null_count==0 and aligned['superclass'].null_count==0
    save(exp/'audit/raw-column-alignment.json',fields);assert all(x['exact'] for x in fields),'Data join mismatch'
    assert column_equal(md['neurotransmitter'],md['nt_consensus_nt'])
    meta=md.to_pandas();mapped=pd.read_parquet(cfg['MB_mapping']).source_neuron_idx.unique();role=roles(meta,mapped)
    role.to_parquet(exp/'roles/neuron-evidence-axes.parquet',index=False)
    save(exp/'regression/classifier-tests.json',self_test())
    meta[meta.annotation_type.eq('DPM')].to_parquet(exp/'audit/DPM-complete-metadata.parquet',index=False)
    evidence=meta[['neuron_idx','body_id','annotation_class','annotation_type','neurotransmitter']+[c for c in meta if c.startswith('nt_')]].copy()
    evidence['annotation_DAN']=role.annotation_DAN;evidence['metadata_dopamine']=role.metadata_dopamine
    evidence[evidence.annotation_DAN!=evidence.metadata_dopamine].to_parquet(exp/'audit/DAN-NT-disagreements.parquet',index=False)
    gt=role.has_raw_ground_truth.to_numpy();same=meta.nt_ground_truth.eq(meta.neurotransmitter)
    conf=meta.nt_predicted_nt_confidence
    summary={'neurons':len(meta),'raw_column_checks':len(fields),'raw_to_metadata_exact':True,'consensus_to_runtime_label_exact':True,
        'ground_truth_present':int(gt.sum()),'ground_truth_absent':int((~gt).sum()),'ground_truth_vs_consensus_mismatch':int((gt & ~same).sum()),
        'ground_truth_agreement_not_independent_accuracy':True,'confidence_null':int(conf.isna().sum()),
        'confidence_outside_0_1':int(((conf<0)|(conf>1)).sum()),'confidence_probability_calibration':'not established; preserve raw values',
        'metadata_dopamine':int(role.metadata_dopamine.sum()),'annotation_DAN':int(role.annotation_DAN.sum()),
        'metadata_dopamine_not_annotation_DAN':int((role.metadata_dopamine & ~role.annotation_DAN).sum()),
        'annotation_DAN_not_metadata_dopamine':int((role.annotation_DAN & ~role.metadata_dopamine).sum()),
        'DPM_count':int(role.annotation_DPM.sum()),'identity_conflicts':int(role.identity_conflict.sum()),
        'DPM_with_raw_ground_truth':int((role.annotation_DPM & role.has_raw_ground_truth).sum()),
        'role_groups':role.role_group.value_counts().to_dict(),'compartment_mapping_unchanged':True}
    save(exp/'audit/metadata-summary.json',summary)
    pcfg=yaml.safe_load((Path(cfg['runtime_root'])/'phase4p.yaml').read_text())
    policy=pcfg['sign_policy'];sign=[]
    for value in meta.neurotransmitter:
        key='' if pd.isna(value) else str(value).strip().lower()
        sign.append(policy['unknown_sign'] if not key else policy['unclear_sign'] if key=='unclear' else policy['neurotransmitter_sign'][key])
    sign=np.asarray(sign,np.int8);saved=np.load(cfg['saved_sign']);assert np.array_equal(sign,saved)
    np.save(exp/'audit/reconstructed-sign.int8.npy',sign)
    rt=phase4p.CanonicalRuntime(pcfg,Path(cfg['runtime_root']),exp);sim=DeterministicAPLSim(rt)
    expected=rt.raw_np.astype(np.float32)*sign[sim.sources]*float(pcfg['model']['w_syn_mv_per_anatomical_synapse'])
    actual=sim.base.cpu().numpy();assert np.array_equal(expected,actual)
    save(exp/'regression/sign-weight-chain.json',{'neurons':len(sign),'edges':len(expected),'saved_sign_exact':True,
        'runtime_effective_weights_exact':True,'base_weight_changes':0,'sign_changes':0,'raw_synapse_sum':int(rt.raw_np.sum()),
        'sign_policy':policy,'fast_current_sign_is_model_assumption_not_receptor_measurement':True})
    wc=yaml.safe_load((Path(cfg['prior_model'])/'phase4w.yaml').read_text());watch=wc['watch_neuron_idx'];rows=[];reg=[];groupdrive=[]
    start=time.perf_counter();oldroot=Path(cfg['prior_root'])
    for oi in wc['odor_indices']:
        name=f'{oi}-{wc["seed"]}-offset';arr=pd.read_parquet(oldroot/'drive'/f'{name}-actual-edge-arrivals.parquet')
        src=arr.source_neuron_idx.to_numpy(np.int32)
        arr['role_group']=role.role_group.to_numpy()[src]
        arr['metadata_dopamine']=role.metadata_dopamine.to_numpy()[src]
        arr['annotation_DAN']=role.annotation_DAN.to_numpy()[src];arr['annotation_DPM']=role.annotation_DPM.to_numpy()[src]
        arr.to_parquet(exp/'roles'/f'{name}-reclassified-arrivals.parquet',index=False)
        for (target,group),g in arr[arr.arrival_step>=5000].groupby(['target_neuron_idx','role_group']):
            w=g.effective_weight.to_numpy(float);groupdrive.append({'trial':name,'neuron_idx':int(target),'role_group':group,
                'positive_sum':float(w[w>0].sum()),'negative_sum':float(w[w<0].sum()),'edge_events':len(g)})
        with np.load(oldroot/'states'/f'{name}-watched.npz') as z:st=z['state']
        for wi,node in enumerate(watch):
            a=arr[arr.target_neuron_idx==node];by={int(t):g for t,g in a.groupby('arrival_step',sort=False)}
            with np.load(oldroot/'replay'/f'{name}-{node}-base.npz') as b:
                # Last-step overlay has no effect on recorded post-state V/g/events.
                base={k:b[k].copy() for k in b.files}
            overlay=np.r_[st[1:,1,wi].astype(float)-st[:-1,5,wi].astype(float),0.]
            for mode in cfg['replay_modes']:
                ar=np.zeros(len(st),np.float32)
                for t,g in by.items():
                    omit=np.zeros(len(g),bool)
                    if mode=='omit_DPM_annotation':omit=g.annotation_DPM.to_numpy()
                    elif mode=='omit_DAN_annotation':omit=g.annotation_DAN.to_numpy()
                    elif mode=='omit_dopamine_without_DAN_or_DPM':omit=(g.metadata_dopamine & ~g.annotation_DAN & ~g.annotation_DPM).to_numpy()
                    elif mode=='omit_legacy_metadata_dopamine':omit=g.metadata_dopamine.to_numpy()
                    w=g.loc[~omit,'effective_weight'].to_numpy(np.float32)
                    if len(w):ar[t]=np.cumsum(w,dtype=np.float32)[-1]
                rr=phase4w.replay(ar,overlay,rt)
                if mode in ('base','omit_legacy_metadata_dopamine'):
                    oldmode='base' if mode=='base' else 'omit_DAN'
                    with np.load(oldroot/'replay'/f'{name}-{node}-{oldmode}.npz') as old:
                        eq={k:bool(np.array_equal(old[k],v)) for k,v in zip(('event_steps','v','g'),rr[:3])}
                    assert all(eq.values());reg.append({'trial':name,'neuron_idx':node,'mode':mode,**eq})
                rows.append({'trial':name,'neuron_idx':node,'mode':mode,'spikes':len(rr[0]),'late_spikes':int((rr[0]>=5000).sum())})
                np.savez_compressed(exp/'replay'/f'{name}-{node}-{mode}.npz',event_steps=rr[0],v=rr[1],g=rr[2])
        print(json.dumps({'odor_complete':oi,'CPU_replays':len(rows)}),flush=True)
    pd.DataFrame(rows).to_parquet(exp/'replay/results.parquet',index=False)
    pd.DataFrame(groupdrive).to_parquet(exp/'roles/reclassified-drive.parquet',index=False)
    save(exp/'regression/replay-compatibility.json',reg)
    assert phase4s.evidence_regression(cfg,exp,rt,sim)
    save(exp/'audit/results.json',{'phase':'4X','status':'AUDIT-CLASSIFICATION-REPAIR-PASS','metadata':summary,
        'raw_import_bug_found':False,'audit_role_conflation_fixed':True,'replays':len(rows),'compatibility_checks':len(reg),
        'reference_pass':True,'sign_weight_chain_exact':True,'frozen_brain_modified':False,'learning_candidate':False,
        'biology_identity_issue_resolved':False,'limit':'conditional replay only; metadata DAN is annotation identity not ground-truth certification'})
    save(exp/'performance/summary.json',{'CPU_replay_and_reference_seconds':time.perf_counter()-start,'GPU':torch.cuda.get_device_name(0)})

def seal(cfg,exp,model):
    pre=json.loads((exp/'protocol/preregistered.json').read_text());now=checks(cfg);hist=phase4p_finalize.integrity()
    ok=now==json.loads((exp/'audit/preflight.json').read_text()) and hist['all_static_checks_passed']
    ok &= all(digest(model/n)==h for n,h in pre['sources'].items()) and all(digest(Path(v['path']))==v['sha256'] for v in pre['inputs'].values())
    save(exp/'audit/frozen-integrity.json',{'pass':ok,'history':hist,'current':now});assert ok
    r=json.loads((exp/'audit/results.json').read_text());r['integrity_pass']=True;save(exp/'audit/results.json',r)
    df=pd.read_parquet(exp/'replay/results.parquet');lines=['# Phase 4X — transmitter provenance and audit-role repair','',
        'Status: AUDIT-CLASSIFICATION-REPAIR-PASS. No frozen brain/sign/raw changes; no learning candidate.','',
        'All original annotation and neurotransmitter columns match raw by body_id. Runtime labels equal raw consensus_nt; reconstructed sign and effective weights match frozen runtime.',
        'DPM has raw predicted/consensus dopamine with no ground_truth. Its confidence is a supplied field, not established calibration probability.',
        'Repair: separate transmitter metadata, exact DAN annotation class, exact DPM annotation type and frozen compartment-mapping evidence. Do not infer reward valence or DAN identity from transmitter/name alone.','',
        '| KC neuron_idx | conditional replay | late spikes range (6 known odors) |','|---:|---|---:|']
    for (node,mode),g in df.groupby(['neuron_idx','mode']):lines.append(f'| {node} | {mode} | {g.late_spikes.min()}–{g.late_spikes.max()} |')
    lines+=['','Six old odors, one fixed seed per odor; no new blind generalization claim. Other recorded arrivals and APL remain fixed in isolated replay.',
        'The last APL subtraction is not applied in CPU comparisons because stored V/g are post-state, pre-APL; it cannot affect any recorded event/state and no final post-APL g comparison is claimed.',
        'Consensus/ground_truth agreement is a provenance check, not independent prediction accuracy. No confidence threshold, label correction, sign change or gain change is inferred.']
    (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']))