"""Exactly two registered evaluation-input batches; no brain, US or learning."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path
import sys
import time

import numpy as np

import boundary_plan as PLAN
import execution_authority as A

MODEL=Path(__file__).resolve().parent
LAB=MODEL.parent/'olfactory-conditioning-lab-v1'
PINS={
    'conditioning_inputs.py':'19382e066e450140bb653a3661f980ffe7cc62f412020ef57df0b6e6b3cbb8ba',
    'conditioning_preflight_inputs.py':'846d2141b472f18b1bb728eed4aa33876896d4b33dece364787c81204f62995f',
    'evaluation_plan.py':'991c813ff3322b230947f302b12b04ec382c3b632817f94b698101e4b17efa7d',
    'trial_plan.py':'6fed3f494d075339a7e5b001d43b2598640b189e038952fbd4d41200bacac587',
    'seed_registry.py':'0da033b132095ec3087b5d9f64e53c2e391fa50ab0b7653b31296eeb84799b7e',
}


def save(path,value):
    with Path(path).open('x',encoding='utf-8') as f:json.dump(value,f,sort_keys=True,indent=2,allow_nan=False)


def run(authority_file,authority_sha,out_dir):
    A.require(sys.dont_write_bytecode and A.digest(authority_file)==authority_sha,'Use -B with external registration pin')
    reg=A.read(authority_file)
    A.require(reg['schema']=='conditioning-execution-registration-v1','Wrong input generation authority')
    PLAN.validate_plan(reg['plan'])
    for path,sha in {**reg['source_sha256'],**reg['input_sha256'],**reg['prepared_file_sha256']}.items():
        A.require(A.pin(path)['sha256']==sha,'Sampler source/input drift')
    for name,sha in PINS.items():
        A.require(A.digest(LAB/name)==sha and reg['input_sha256'].get(str(LAB/name))==sha,'Unpinned frozen input dependency: '+name)
    A.require(not any(name in sys.modules for name in ('conditioning_inputs','conditioning_preflight_inputs','evaluation_plan')),
              'Sampler needs fresh scientific imports')
    sys.path.insert(0,str(LAB))
    import conditioning_preflight_inputs as I
    import conditioning_inputs as CI
    import evaluation_plan as EP
    rates=I.load_rate_tables('/data/digital-drosophila')
    full=reg['original_full_plan'];A.require(A.digest(full['path'])==full['sha256'],'Full seed/trial plan changed')
    streams=EP._streams(A.read(full['path']))
    exp=Path(authority_file).parent.parent;out=Path(out_dir)
    A.require(not out.exists() and out.parent.is_dir() and out.is_relative_to(exp),'Exclusive new sampler output required')
    out.mkdir();rows=[]
    for cue in ('A','B'):
        declared=reg['plan']['evaluation_streams'][cue];actual=streams[(PLAN.BLOCK,cue)]
        for field in ('seed','seed_bits','stream_id','derivation_sha256'):
            A.require(declared[field]==actual[field],'Evaluation semantic stream differs')
        A.require(declared['key']==actual['semantic_key'] and actual['split']=='evaluation','Evaluation stream provenance differs')
        table=rates['tables'][cue]
        with np.load(exp/f'inputs/static/{cue}-rates.npz',allow_pickle=False) as saved:
            A.require(all(saved[k].dtype==table[k].dtype and saved[k].shape==table[k].shape and
                          saved[k].tobytes()==table[k].tobytes() for k in ('targets','rates_hz')),'Registered original PN mapping changed')
        started=time.perf_counter()
        result=CI.generate_cs_requests(table['targets'],table['rates_hz'],seed=declared['seed'],steps=1200,
            cs_start_tick=0,cs_stop_tick=1000,device='cuda:0',max_uniform_draws=1200*len(table['targets']))
        with (out/f'{cue}-requests.npz').open('xb') as f:np.savez_compressed(f,**result['requests'])
        save(out/f'{cue}-sampler.json',result['audit'])
        rows.append(dict(cue=cue,stream=actual,requests=len(result['requests']['tick']),US_requests=0,
            uniform_draws=1200*len(table['targets']),wall_seconds=time.perf_counter()-started,
            saved_requests_sha256=A.digest(out/f'{cue}-requests.npz')))
    pins={p.name:A.digest(p) for p in out.iterdir() if p.is_file()}
    summary=dict(schema='conditioning-boundary-sampler-v1',complete=True,registration_sha256=authority_sha,
        uniform_batch_calls=2,uniform_draws=sum(r['uniform_draws'] for r in rows),rows=rows,files_sha256=pins,
        brain_executed=False,learning_executed=False,US_added=False,original_rate_mapping_preserved=True,
        inputs_sampled_once=True,planned_retention_reuse_checkpoints_seconds=[0,1,10],
        reuse_executed_here=False,seed_search=False,learning_success_claimed=False)
    save(out/'sampler-receipt.json',summary)
    return summary