"""Exclusive CPU preparation for the finite AO conditioning-interface gate.

This is a two-epoch installation test, not acquisition, retention or a learning
candidate. Source authority and the complete 12-call plan precede CUDA sampling.
Generated requests and learned snapshots receive separate execution receipts;
their not-yet-existing hashes are never invented in the registration.
"""
from __future__ import annotations

import json
from pathlib import Path
import shutil
import subprocess
import sys

import numpy as np

import engineering_prepare as P
import runtime_light as R
import biological_contract as B
import conditioning_preflight_inputs as I

ROOT = Path('/data/digital-drosophila')
MODEL = Path(__file__).resolve().parent
EXP = ROOT / 'experiments/phase4ao-conditioning-interface'
AN = ROOT / 'experiments/phase4an-engineering-validation'
AN_REG = '7317441126f191ccc6e6d28438870b3921537ebe282a6acb163716a690057945'
AN_MANIFEST = 'e8952a48cd3ea6a93c811a1183a1f7b79e19646299924e7688e5f393737a9aa2'
AN_FINAL = '16630d0d666a2467c56b647d021cbdb627a553795f6afcd7f7eb496a2737d815'
PREP_COMMIT = 'c187d7d5a6047f183467e3d860e2529575dcda97'
DRAFT_INVENTORY_SHA = '2c72d47d09b4ad1949dc6cfdc170f64847ff75f95ded5d20b421fb5a75e0eb84'
BUDGET = dict(output_budget_bytes=4 * 1024**3, minimum_free_reserve_bytes=10 * 1024**3,
              max_uniform_draws=2_000_000, child_timeout_seconds=1800)
require, read, save, digest = P.require, P.read, P.save, P.digest


def same(a, b):
    return json.dumps(a, sort_keys=True, allow_nan=False) == json.dumps(b, sort_keys=True, allow_nan=False)


def verify_pins(pins):
    for name, sha in pins.items():
        path = Path(name)
        require(path.is_absolute() and path.is_file() and not path.is_symlink() and
                digest(path) == sha, 'Pinned file changed: ' + name)


def save_npz(path, **arrays):
    with Path(path).open('xb') as stream:
        np.savez_compressed(stream, **arrays)


def source_inventory():
    # Explicitly snapshot the current lab's text sources, including this new
    # coordinator and its tests. Existing 21/20 files are never rewritten.
    paths = sorted([*MODEL.glob('*.py'), *MODEL.glob('*.md')])
    required = ('conditioning_epoch.py', 'conditioning_preflight_inputs.py',
        'conditioning_preflight_compare.py', 'conditioning_preflight_prepare.py',
        'conditioning_preflight_workflow.py', 'test_conditioning_preflight_prepare.py',
        'CONDITIONING_INTERFACE_PROTOCOL.md')
    require(all(MODEL / n in paths for n in required), 'Missing complete coordinator/test/protocol source')
    require(not any(p.is_symlink() for p in paths), 'Linked source is forbidden')
    return {str(p): digest(p) for p in paths}


def plan_descriptor(full):
    return {k: v for k, v in full.items() if k not in ('full_trial_plan', 'payload_sha256')}


def prepare():
    require(sys.dont_write_bytecode, 'Use scientific Python -B')
    require(EXP.is_dir() and not any(EXP.iterdir()) and not EXP.is_symlink(),
            'Preparation requires the explicitly created empty AO directory; never overwrite')
    require(digest(AN / 'manifest.json') == AN_MANIFEST and digest(AN / 'audit/final-status.json') == AN_FINAL,
            'Externally anchored AN controls changed')
    an = P.read_registration(AN_REG)
    helper = P.load_module('phase4ao_original_seal_helper', ROOT / P.AI_HELPER_REL, P.AI_HELPER_SHA)
    an_seal = helper.verify_seal(AN)
    am_seal = P.verify_authority(an['anchors'])
    old = P.verify_old_integrity(an)
    full = I.build_preflight_plan()
    rates = I.load_rate_tables(ROOT)
    populations, biology = B.load_bound_tables(ROOT, an['input_sha256'])
    pins = dict(an['input_sha256'])
    pins.update({str(ROOT / k): v for k, v in I.SOURCE_PINS.items()})
    pins.update(biology['input_sha256'])
    pins.update({str(AN / 'manifest.json'): AN_MANIFEST, str(AN / 'audit/final-status.json'): AN_FINAL,
                 str(AN / 'protocol/preregistered.json'): AN_REG})
    bridge = ROOT / 'models/male-cns-lif/wholebrain-plastic-interface-v1/eligibility_bridge.py'
    phasej = ROOT / 'models/male-cns-lif/mb-physiology-v1/phase4j.py'
    pins[str(bridge)] = '174fb8d14619189e74b095e77a6db18b14c93046dc431c0b55710a5aa2dd2e61'
    pins[str(phasej)] = 'abb4574da85f00bd9ef61024c65993ef81d42651ff8d8d63c15dd69c6222471e'
    original_path = Path(an['selected_case']['original_dir']) / 'spec.json'
    require(str(original_path) in pins, 'Original graph spec not anchored by AN')
    original = read(original_path)
    previous = MODEL / 'draft-source-inventory.json'
    require(digest(previous) == DRAFT_INVENTORY_SHA, 'Previously delivered 20-file source inventory changed')
    inventory = read(previous)
    require(len(inventory['files']) == 20, 'Previous scientific snapshot must contain 20 files')
    pins[str(previous)] = DRAFT_INVENTORY_SHA
    for entry in inventory['files']:
        require(Path(entry['path']).name == entry['path'], 'Nonlocal scientific source inventory entry')
        pins[str(MODEL / entry['path'])] = entry['sha256']
    zero = an['overlays']['zero']
    _, overlay, metadata, _ = R.load_overlay(zero['overlay_file'], zero['overlay_manifest'])
    require(not np.any(overlay['plastic_delta_mv'].view(np.uint32)), 'Initial delta must be exact positive zero')
    kc = np.load(Path(original['apl_bundle_dir']) / 'kc_globals.npy', allow_pickle=False)
    require(np.array_equal(kc, populations['Kenyon_Cell']), 'Canonical metadata and APL KC memberships disagree')
    import pandas as pd
    candidate = pd.read_parquet(ROOT / 'experiments/phase4h-dopamine-plasticity/plastic-mask/plastic-edge-candidates.parquet')
    candidate = candidate.set_index('edge_id').loc[overlay['edge_id']].reset_index()
    require(candidate.edge_id.to_numpy().tobytes() == overlay['edge_id'].tobytes() and
            candidate.target_compartment_verbatim.notna().all(), 'Complete mask compartment mapping is not exact')
    verify_pins(pins)
    sources = source_inventory()
    code_commit = subprocess.run(['git', '-C', str(ROOT), 'rev-parse', 'HEAD'], check=True,
                                 capture_output=True, text=True).stdout.strip()
    require(P.sha(code_commit + '0' * 24), 'Invalid Git code version')
    for path, sha in sources.items():
        relative = Path(path).relative_to(ROOT).as_posix()
        blob = subprocess.run(['git', '-C', str(ROOT), 'show', code_commit + ':' + relative],
                              check=True, capture_output=True).stdout
        require(__import__('hashlib').sha256(blob).hexdigest() == sha, 'Registered source is not the committed Git blob: ' + relative)
    for name in ('protocol/source-snapshot', 'inputs/static', 'cases', 'runs', 'updates', 'receipts', 'logs', 'audit', 'regression'):
        (EXP / name).mkdir(parents=True)
    for path in sources:
        shutil.copyfile(path, EXP / 'protocol/source-snapshot' / Path(path).name)
    save(EXP / 'regression/AN-seal-before.json', an_seal)
    save(EXP / 'regression/AM-seals-before.json', am_seal)
    save(EXP / 'regression/old-integrity-before.json', old)
    save(EXP / 'inputs/static/biology-audit.json', biology)
    save(EXP / 'inputs/static/rates-audit.json', rates['audit'])
    save_npz(EXP / 'inputs/static/populations.npz', **populations)
    save_npz(EXP / 'inputs/static/candidate-map.npz', edge_id=candidate.edge_id.to_numpy(np.int64),
        source_neuron_idx=candidate.source_neuron_idx.to_numpy(np.int32),
        target_neuron_idx=candidate.target_neuron_idx.to_numpy(np.int32),
        target_compartment_verbatim=candidate.target_compartment_verbatim.to_numpy(dtype=str))
    table_metadata = {}
    for cue, table in rates['tables'].items():
        save_npz(EXP / f'inputs/static/{cue}-rates.npz', targets=table['targets'], rates_hz=table['rates_hz'])
        table_metadata[cue] = {k: v for k, v in table.items() if k not in ('targets', 'rates_hz')}
    save(EXP / 'inputs/static/rate-tables.json', table_metadata)
    save(EXP / 'protocol/full-trial-plan.json', full['full_trial_plan'])
    descriptor = plan_descriptor(full)
    cases = {}
    reset_seed = original['seed']
    for row in descriptor['calls']:
        epoch = descriptor['selected_epochs'][row['input_epoch_index']]
        case = EXP / 'cases' / row['row_id']
        case.mkdir()
        spec, derivation = I.derive_case_spec(original, row, epoch, rates['tables'][epoch['cs']['cue']],
            brain_reset_seed=reset_seed, budget_root=str(EXP), output_budget_bytes=BUDGET['output_budget_bytes'])
        R.validate_spec(spec, case)
        save(case / 'spec.json', spec)
        save(case / 'derivation.json', derivation)
        cases[row['row_id']] = str(case)
    derived = {str(p): digest(p) for parent in ('inputs/static', 'cases', 'protocol/source-snapshot', 'regression')
               for p in (EXP / parent).rglob('*') if p.is_file()}
    derived[str(EXP / 'protocol/full-trial-plan.json')] = digest(EXP / 'protocol/full-trial-plan.json')
    reg = dict(schema='phase4ao-conditioning-interface-registration-v1', plan=descriptor,
        full_plan_payload_sha256=full['payload_sha256'], full_trial_plan_file_sha256=derived[str(EXP / 'protocol/full-trial-plan.json')],
        source_sha256=sources, input_sha256=pins, prepared_file_sha256=derived, cases=cases,
        initial_overlay=zero, original_spec=str(original_path), brain_reset_seed=reset_seed,
        bridge_source=str(bridge), phase4j_source=str(phasej), budget=BUDGET,
        git_parent_commit=PREP_COMMIT, git_code_commit=code_commit, committed_source_blobs_verified=len(sources), AN_registration_sha256=AN_REG,
        generation=dict(device='cuda:0', GPU_UUID='GPU-79f3a85d-adce-7833-7003-6488140ca596',
            conflict_policy='error', highbit_probe_seed=descriptor['selected_epochs'][0]['cs']['seed'],
            highbit_comparison='full63bit versus low32bit, one fixed 1200xPN matrix each; never a seed search',
            original_L_sampling_AST='only rate tensor, private generator and full-shape forced-spike expression',
            trace_backend='torch', actual_input_files='created only after external registration SHA recorded'),
        scope=dict(epoch_boundary_updates=True, within_trial_updates=False, full_conditioning=False,
            candidate_generated=False, learning_success_claimed=False, neural_state_carried=False,
            delta_carried=True, independent_model_branch=True),
        success_gate='12 complete calls; exact reference/light events and final state/queues; exact D0/frozen identity; real D1 delta installed in epoch 2; nonzero D1 update coverage required',
        missing_nonzero_coverage='INCOMPLETE, do not choose another input',
        failure_policy='preserve all partials; no automatic run retry or result-selected seed',
        retention_supported=False, future_hashes_fabricated=False, new_brain_runs_during_preparation=0)
    verify_pins({**pins, **sources, **derived})
    save(EXP / 'protocol/preregistered.json', reg)
    return digest(EXP / 'protocol/preregistered.json')


def read_registration(expected):
    path = EXP / 'protocol/preregistered.json'
    require(P.sha(expected) and digest(path) == expected, 'Registration differs from external anchor')
    reg = read(path)
    require(reg['schema'] == 'phase4ao-conditioning-interface-registration-v1' and
            reg['AN_registration_sha256'] == AN_REG and reg['budget'] == BUDGET and
            reg['plan']['total_calls'] == 12 and reg['plan']['total_ticks'] == 14400,
            'Unexpected finite registration')
    verify_pins({**reg['source_sha256'], **reg['input_sha256'], **reg['prepared_file_sha256']})
    return reg


if __name__ == '__main__':
    require(sys.argv[1:] == ['prepare'], 'Only exclusive prepare is supported')
    print('REGISTRATION_SHA256', prepare(), flush=True)