"""Register only normal acquisition; never sample inputs or execute a brain.

Prior sealed files stay read-only. Complete historical integrity is checked at
the phase boundary; each trial rehashes its actual runtime dependency closure,
not every historical engineering result. No source or budget changes after
registration. Outcome files have no invented pre-outcome hashes.
"""
from __future__ import annotations
import argparse
import copy
from datetime import datetime, timezone, timedelta
import shutil
import subprocess
import sys
from normal_io import *
import normal_plan as PLAN

AQ_EXP = ROOT/'experiments/phase4aq-conditioning-execution'
AP_EXP = ROOT/'experiments/phase4ap-conditioning-interface-v2'
AO_EXP = ROOT/'experiments/phase4ao-conditioning-interface'
AQ_ANCHORS = {
    'protocol/preregistered.json':'dad0ba33ded8c6725662c419525ef4185a6ff169042a402c18d0bd760dad7358',
    'manifest.json':'2dc0feca75f718f2ee0d7c9240da503221823573fd25f1344d2fd97d97eb2db7',
    'audit/final-status.json':'dcbb58f296bd0d51e3c8757f0c45dd5585b3a76f8c2a0b5881f4c14076721882',
}
H_ANCHORS = {
    'audit/phase4h.py':'28bb2ab7c4c95af3d95b0bf0bc45a5a9bea791996c59ce296ec66b71e01b1695',
    'mappings/mbon-to-downstream-actual-edges.parquet':'76bac27c50cec41c48fd6414bcef5a6e270753ac4062b92537f38226018e9278',
    'manifests/checksums.sha256':'1ea0394fd0477bb40d656ed0854108964f3809d633b5561bed9de43cda630ba1',
    'manifests/manifest.json':'d1f187f164f7775a1a02fff259c8f5e31ae035929d808c3876cfc0013d19c73c',
}

# Exact arrays consumed by the pinned isolated worker, not every reference file
# colocated in its graph/APL directories. Tests compare these to that worker.
GRAPH_FILES = ('nodes.npy', 'edge_id.npy', 'src.npy', 'dst.npy', 'signed_raw.npy', 'ptr.npy', 'weights.npy')
APL_FILES = ('input_matrix.npy', 'output_matrix.npy', 'kc_globals.npy', 'apl_globals.npy')

def runtime_array_paths(template):
    return {str(Path(template['graph_dir'])/name) for name in GRAPH_FILES} | {
        str(Path(template['apl_bundle_dir'])/name) for name in APL_FILES} | {template['coefficients_path']}

def readout_populations(population_file):
    import numpy as np
    import pyarrow.parquet as pq
    h=ROOT/'experiments/phase4h-dopamine-plasticity'
    verify({str(h/p):s for p,s in H_ANCHORS.items()})
    rows=pq.read_table(h/'mappings/mbon-to-downstream-actual-edges.parquet',columns=['target_neuron_idx'])['target_neuron_idx'].to_numpy()
    require(rows.dtype==np.int32 and len(rows)==39279 and np.all((rows>=0)&(rows<166700)), 'Original H endpoint schema differs')
    old=npz(population_file)
    result={k:old[k] for k in ('MBON','MBON01','MBON12')}
    result['legacyDownstream']=np.unique(rows).astype(np.int32)
    result['nonMBONDownstream']=np.setdiff1d(result['legacyDownstream'],result['MBON'])
    require(not len(np.setdiff1d(result['MBON'],result['legacyDownstream'])),'H downstream no longer includes all MBON')
    hashes={k:hashlib.sha256(v.tobytes()).hexdigest() for k,v in result.items()}
    require(hashes['legacyDownstream']=='ce9edbb56854be8f6891bfb83559dd7529c70acac6867864947f32f933175dbc' and
        hashes['nonMBONDownstream']=='01e3fbdde8bf2558665f142cdf546d8b89b4898427b90e2f4368f21f4e241a7f',
        'Actual H readout membership changed')
    return result,dict(source_sha256={str(h/p):s for p,s in H_ANCHORS.items()},
        original_populations=pin(population_file),definition='unique actual H MBON-edge target_neuron_idx, int32; no plastic/active subset',
        nonMBON_definition='exact set difference legacyDownstream minus full anatomical MBON',
        counts={k:len(v) for k,v in result.items()},raw_membership_sha256=hashes,
        historical_current_model_source_used=False,original_sealed_audit_source_used=True)

def historical_integrity():
    verify({str(AQ_EXP/p):s for p,s in AQ_ANCHORS.items()})
    ar = read(AQ_EXP/'protocol/preregistered.json')
    verify(ar['source_sha256'])
    # This function is executed in a separate preparation/integrity CPU process.
    for p in (AQ, AP, LAB): sys.path.insert(0,str(p))
    old = load_source('_normal_previous_integrity',AQ/'prepare_execution.py',ar['source_sha256'][str(AQ/'prepare_execution.py')])
    prior = old.integrity()
    helper = load_source('_normal_AQ_seal_reader', MODEL.parent/'pn-recruitment-audit-v1/publication.py',
        '7a79c499098b74163e7bb16adabcbc54e7069dafa204d0b99c7f4c03d79f8016')
    sealed = helper.verify_seal(AQ_EXP)
    require(read(AQ_EXP/'audit/final-status.json')['status'] ==
        'PASS-EXACT-UPDATED-BOUNDARY-ENGINEERING-NOT-LEARNING','AQ gate differs')
    return dict(complete=True,prior=prior,AQ=sealed,old_experiments_reexecuted=False,
        frozen_five_reexecuted=False,scope='actual historical seal bytes reverified, preserved numerical evidence; not a new numerical trial')

def resource_policy():
    """Engineering ceiling, not a measured all-condition forecast or stop-for-success."""
    return dict(root=str(EXP),output_cap_bytes=256*1024**3,minimum_free_reserve_bytes=100*1024**3,
        child_timeout_seconds=1800,workflow_timeout_seconds=21*86400,max_uniform_draws=10_000_000,
        next_child_headroom_bytes=4*1024**3,
        execution='serial_single_4090',no_automatic_retry=True,
        forecast_scope='rough planning arithmetic from one AQ warm training cycle and one probe; not condition-complete measurement',
        AQ_measured_training_cycle_seconds=24.771224757998425,AQ_measured_probe_seconds=10.394962482001574,
        AQ_linear_wall_scenario_seconds=30000*24.771224757998425+3920*10.394962482001574,
        uncertainties=['startup and materialization','event-dependent output size','all control/DAN/odor activity',
            'ledger growth','shared service contention','final full integrity audit'],
        cap_or_timeout_policy='preserve all outputs, stop incomplete; no shortening epochs or selecting favorable outcomes')

def prepare():
    require(sys.dont_write_bytecode and unlinked(EXP).is_dir() and not any(EXP.iterdir()), 'Exact new empty EXP and -B required')
    files=sorted(MODEL.iterdir())
    require(files and all(p.is_file() and not p.is_symlink() and p.suffix in ('.py','.md') for p in files), 'Source-only model directory')
    sources={str(p):digest(p) for p in files}
    commit=subprocess.check_output(['git','-C',str(ROOT),'rev-parse','HEAD']).decode().strip()
    for path,sha in sources.items():
        blob=subprocess.check_output(['git','-C',str(ROOT),'show',commit+':'+Path(path).relative_to(ROOT).as_posix()])
        require(hashlib.sha256(blob).hexdigest()==sha, 'Uncommitted source '+path)
    budget=resource_policy()
    budget['registered_utc']=datetime.now(timezone.utc).isoformat()
    budget['not_after_utc']=(datetime.now(timezone.utc)+timedelta(seconds=budget['workflow_timeout_seconds'])).isoformat()
    require(shutil.disk_usage(EXP).free >= budget['output_cap_bytes']+budget['minimum_free_reserve_bytes'], 'Insufficient reserved capacity')
    for name in ('protocol/source-snapshot','inputs/static','jobs','cases','runs','updates','validation','results','receipts','logs','audit','regression'):
        unlinked(EXP/name).mkdir(parents=True)
    tests=subprocess.run([sys.executable,'-B',str(MODEL/'run_cpu_tests.py'),'--output',str(EXP/'audit/cpu-tests.json')],
        capture_output=True,text=True,timeout=1800,env={**dict(__import__('os').environ),'CUDA_VISIBLE_DEVICES':''})
    save(EXP/'audit/cpu-tests.process.json',dict(exit_code=tests.returncode,stdout=tests.stdout,stderr=tests.stderr,
        actual_parent_observation=True))
    require(tests.returncode==0 and read(EXP/'audit/cpu-tests.json')['pass_all'] is True,
        'CPU/source tests failed; no sampler or brain started')
    # Old integrity readers may own a private R/C module graph. Keep them in an
    # exec-isolated CPU process, never the table/registration interpreter.
    integrity_command=[sys.executable,'-B',str(MODEL/'normal_prepare.py'),'integrity',
        '--output',str(EXP/'regression/integrity-before.json')]
    checked=subprocess.run(integrity_command,capture_output=True,text=True,timeout=1800,
        env={**dict(__import__('os').environ),'CUDA_VISIBLE_DEVICES':''})
    save(EXP/'regression/integrity-before.process.json',dict(argv=integrity_command,exit_code=checked.returncode,
        stdout=checked.stdout,stderr=checked.stderr,actual_parent_observation=True))
    require(checked.returncode==0,'Historical integrity child failed; preserve artifacts')
    before=read(EXP/'regression/integrity-before.json');require(before['complete'] is True,'Integrity child incomplete')
    aq=read(AQ_EXP/'protocol/preregistered.json')
    full=read(PLAN.AO_TRIAL_PATH)
    require(digest(PLAN.AO_TRIAL_PATH)==PLAN.AO_TRIAL_FILE_SHA256,'Original full plan pin differs')
    evaluation=PLAN.expand_pinned_evaluation(full)
    plan=PLAN.build_normal_plan(full,evaluation)
    from normal_inputs import derive_rates, inventory_from_plan, PINS, dependencies
    rate_tables,rate_audit=derive_rates()
    inputs=inventory_from_plan(plan)
    template_path=AP_EXP/'cases/D1--reference--epoch-001/spec.json'
    authority_path=AP_EXP/'jobs/D1--reference--epoch-001.authority.json'
    template=read(template_path); epoch=read(authority_path)
    pinned={**aq['input_sha256'],**aq['source_sha256'],**aq['prepared_file_sha256']}
    for path in (template_path,authority_path,AO_EXP/'inputs/static/populations.npz'):
        require(str(path) in pinned and digest(path)==pinned[str(path)],'Original template/population not pinned')
    for path in sources:
        with Path(path).open('rb') as i,(EXP/'protocol/source-snapshot'/Path(path).name).open('xb') as o:shutil.copyfileobj(i,o)
    save(EXP/'protocol/normal-plan.json',plan)
    save(EXP/'protocol/evaluation-policy.json',evaluation['config'])
    save(EXP/'protocol/input-inventory.json',inputs)
    save(EXP/'protocol/case-template.json',template)
    # Freeze scientific predicates before any new CS sample or outcome.
    gate_sources={str(AQ/'learning_gates.py'):'8ae9bcb77fe14dbdda164f86ffdbfc1e2976be752766ef526777577223b4dd99',
        str(LAB/'metrics.py'):'47319e21db3c8b030b1bafbf0703cd86fb588d23bb29436accbacdba73d08c3c'}
    verify(gate_sources)
    save(EXP/'protocol/gates.json',dict(source_sha256=gate_sources,all_eight_cells_required=True,
        same_replicate_conjunction_at_least=3,training_replicates=5,terminal_checkpoint=50,
        normal_failure_prevents_followups=True,intermediate_checkpoints_cannot_rescue_terminal=True,
        downstream_policy='nonMBONDownstream-pre-outcome-operationalization-v1',
        H_mapping='each-DAN-control-source-separately-both-pass-v1',
        physiological_and_matched_energy_not_pooled=True,unknown_compartments_not_guessed=True))
    save(EXP/'inputs/static/rate-audit.json',rate_audit)
    for key,table in rate_tables.items():save_npz(EXP/f'inputs/static/{key}.npz',table)
    populations,pop_audit=readout_populations(AO_EXP/'inputs/static/populations.npz')
    save_npz(EXP/'inputs/static/readout-populations.npz',populations)
    save(EXP/'inputs/static/readout-population-binding.json',pop_audit)
    # Actual consumed neural dependency closure, all inherited byte pins.
    required=runtime_array_paths(template)
    require(required <= pinned.keys(), 'Consumed runtime arrays missing inherited pins: '+str(sorted(required-pinned.keys())))
    original_sources={p:s for p,s in epoch['frozen_inputs'].items() if p.startswith(str(ROOT/'models/male-cns-lif/') )
        and ('author-' in p or 'apl-timing-lab-v1/' in p)}
    runtime={p:pinned[p] for p in required}
    runtime.update(original_sources)
    roles={k:v for k,v in epoch['files'].items() if k not in ('spec','requests','state','overlay_before','overlay_manifest_before')}
    runtime.update({v['path']:v['sha256'] for v in roles.values()})
    save(EXP/'protocol/epoch-template.json',dict(schema=epoch['schema'],files=roles,frozen_inputs=runtime))
    static={str(LAB/n):s for n,s in {**PLAN.SOURCE_PINS,**PINS}.items()}
    static.update({str(ROOT/p):s for p,s in rate_audit['original_source_sha256'].items()})
    external={**runtime,**static,**gate_sources,**pop_audit['source_sha256'],**{str(AQ_EXP/p):s for p,s in AQ_ANCHORS.items()},
        str(PLAN.AO_TRIAL_PATH):PLAN.AO_TRIAL_FILE_SHA256,str(AP/'isolated_workers.py'):AP_SHA,
        str(authority_path):pinned[str(authority_path)],str(template_path):pinned[str(template_path)],
        str(AO_EXP/'inputs/static/populations.npz'):pinned[str(AO_EXP/'inputs/static/populations.npz')],
        str(AQ/'validate_update_cpu.py'):aq['source_sha256'][str(AQ/'validate_update_cpu.py')]}
    for name in ('conditioning_epoch.py','conditioning_preflight_compare.py','engineering_compare.py'):
        path=str(LAB/name);require(path in pinned,'Missing original verifier source pin');external[path]=pinned[path]
    zero=aq['initial_overlay']
    external.update({p:digest(p) for p in zero.values()})
    verify(external)
    prepared={p:row['sha256'] for p,row in inventory(EXP).items()}
    reg=dict(schema='normal-conditioning-registration-v1',scientific_scope='normal-acquisition-only',
        subsequent_assays_authorized=False,source_code_commit=commit,source_sha256=sources,input_sha256=external,
        static_input_sha256=static,runtime_input_sha256=runtime,prepared_sha256=prepared,budget=budget,
        initial_overlay=zero,population_file=str(AO_EXP/'inputs/static/populations.npz'),
        readout_population_file=str(EXP/'inputs/static/readout-populations.npz'),
        readout_population_binding_file=str(EXP/'inputs/static/readout-population-binding.json'),
        original_full_trial_plan=pin(PLAN.AO_TRIAL_PATH),normal_inventory=plan['inventory'],
        evidence_scope='AQ validates exact interface, not learning; this new registration admits actual normal outcomes',
        brain_parameters_changed=False,learning_success_claimed=False,
        execution_order='original fixed block/arm order; actual A/B baseline each arm; no outcome selection',
        biological_hypotheses='D18 branch; PAM15/PPL103 compartment hypotheses only, not all-PAM/PPL functional assignment',
        historical_integrity_scope=before['scope'],failure_policy='stop preserve partials, never auto retry',
        administrative_supervisor_directory=str(ROOT/'logs/phase4ar-normal-conditioning'),
        supervisor_logging_scope='outside immutable EXP ledger; real top-level exit and logs must be included by final outer auditor',
        full_plan_in_old_frozen_location=True,future_artifact_hashes_fabricated=False)
    save(EXP/'protocol/preregistered.json',reg)
    return digest(EXP/'protocol/preregistered.json')

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('mode',choices=('prepare','integrity'));p.add_argument('--output')
    a=p.parse_args()
    if a.mode=='prepare':print('REGISTRATION_SHA256',prepare(),flush=True)
    else:
        require(a.output,'Exclusive report path required');save(a.output,historical_integrity())