"""Read existing records and replay PN targets only; no GPU or counterfactual."""
from functools import lru_cache
import time
import traceback
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from context import (ROOT, EXP, AH_MODEL, AG, COEFF, GROUPS, read, save, digest,
                     check, bit_equal, load_module)
from pn_observer import PNObserver
from apl_sources import decompose_apl_sources


def first_ticks(event_tick, event_node, targets):
    first = np.full(166700, np.iinfo(np.int32).max, np.int32)
    np.minimum.at(first, event_node, event_tick)
    first[first == np.iinfo(np.int32).max] = -1
    return first[targets]


def temporal_relation(source_tick, target_tick):
    if source_tick < 0 or target_tick < 0:
        return 'unobserved_source' if source_tick < 0 and target_tick >= 0 else (
            'unobserved_target' if source_tick >= 0 else 'both_unobserved')
    return 'before' if source_tick < target_tick else ('same_tick' if source_tick == target_tick else 'after')


def min_observed(values):
    values = np.asarray(values)
    values = values[values >= 0]
    return int(values.min()) if len(values) else -1


class EventSink:
    """Streaming accepted antecedents, never silently truncate."""
    def __init__(self, path):
        ints64 = ('edge_id', 'incoming_edge_row')
        ints32 = ('source_neuron_idx', 'target_neuron_idx', 'source_spike_tick', 'arrival_tick',
                  'weight_sign', 'ref_before', 'ref_at_arrival')
        bools = ('target_is_design', 'accepted', 'available_at_integration')
        floats = ('weight_SI', 'old_v_SI', 'old_g_SI', 'v_integrated_SI', 'v_prethreshold_SI',
                  'threshold_margin_SI', 'g_before_arrivals_SI')
        self.schema = pa.schema([(k, t) for names, t in ((ints64, pa.int64()), (ints32, pa.int32()),
             (bools, pa.bool_()), (floats, pa.float64()), (('source_group', 'record_scope'), pa.string())) for k in names])
        assert not path.exists()
        self.writer = pq.ParquetWriter(path, self.schema, compression='zstd')
        self.rows = 0

    def __call__(self, rows):
        self.writer.write_table(pa.Table.from_pylist(rows, schema=self.schema))
        self.rows += len(rows)

    def close(self):
        self.writer.close()


@lru_cache(maxsize=2)
def graph_view(graph_name):
    core = load_module('sealed_ah_replay_ai_' + graph_name, AH_MODEL / 'replay_core.py')
    folder = AG / 'graphs' / graph_name
    nodes = np.load(folder / 'nodes.npy', mmap_mode='r')
    pn = np.load(AG / 'apl-bundle/pn_globals.npy')
    apl = np.load(AG / 'apl-bundle/apl_globals.npy')
    targets, positions, _ = np.intersect1d(nodes, pn, assume_unique=True, return_indices=True)
    assert len(targets) > 0, 'Empty PN fixture subset: stop and report nonapplicable; never fabricate targets'
    lookup = np.full(len(nodes), -1, np.int32)
    lookup[positions] = np.arange(len(targets), dtype=np.int32)
    dst = np.load(folder / 'dst.npy', mmap_mode='r')
    selected = np.flatnonzero(lookup[dst] >= 0)
    src = nodes[np.load(folder / 'src.npy', mmap_mode='r')[selected]]
    ids = np.load(folder / 'edge_id.npy', mmap_mode='r')[selected]
    weights = np.load(folder / 'weights.npy', mmap_mode='r')[selected].copy()
    weights[np.isin(src, apl)] = 0.0
    raw = np.load(ROOT / 'datasets/male-cns-v1.0/processed/neuron-only/gpu/csr-weights.int64.npy', mmap_mode='r')[ids]
    assert raw.dtype == np.int64 and np.all(raw > 0)
    incoming = core.IncomingIndex(src, lookup[dst[selected]], ids, weights, len(targets))
    groups = pd.read_parquet(EXP / 'protocol/neuron-source-groups.parquet').audit_source_group.to_numpy()
    codes = np.array([GROUPS.index(x) for x in groups], np.int32)[src]
    audit = dict(graph=graph_name, PN_targets=len(targets), incoming_edges=len(ids), contacts=int(raw.sum()),
        zeroed_fast_APL_edges=int(np.isin(src, apl).sum()), graded_APL_to_PN=False,
        groups={g: dict(edges=int((codes == i).sum()), contacts=int(raw[codes == i].sum())) for i, g in enumerate(GROUPS)},
        source_graph_unchanged=True, effective_fast_APL_zero_mask_preserved=True)
    return core, incoming, raw, codes, targets, positions, audit


@lru_cache(maxsize=2)
def pn_kc_edges(graph_name):
    folder = AG / 'graphs' / graph_name
    nodes = np.load(folder / 'nodes.npy', mmap_mode='r')
    pn = np.load(AG / 'apl-bundle/pn_globals.npy')
    kc = np.load(AG / 'apl-bundle/kc_globals.npy')
    src = nodes[np.load(folder / 'src.npy', mmap_mode='r')]
    dst = nodes[np.load(folder / 'dst.npy', mmap_mode='r')]
    picked = np.flatnonzero(np.isin(src, pn) & np.isin(dst, kc))
    return src[picked], np.searchsorted(kc, dst[picked]), np.load(folder / 'weights.npy', mmap_mode='r')[picked]


def pn_kc_before_first(graph_name, state, kc_first, design):
    """All PN edge arrivals strictly before first KC spike; no refractory yet."""
    src, dst, weights = pn_kc_edges(graph_name)
    rows = {'neuron_idx': np.load(AG / 'apl-bundle/kc_globals.npy').copy(), 'first_spike_tick': kc_first.copy()}
    for name in ('design_ALPN', 'non_design_ALPN'):
        rows[name + '_strict_pre_first_edge_events'] = np.zeros(len(kc_first), np.int64)
        rows[name + '_strict_pre_first_positive_V'] = np.zeros(len(kc_first), np.float64)
        rows[name + '_strict_pre_first_negative_abs_V'] = np.zeros(len(kc_first), np.float64)
    # Fixed graph edges only. Before the first spike, initially zero ref has never been reset.
    for source in np.unique(src):
        edges = np.flatnonzero((src == source) & (kc_first[dst] >= 0))
        spikes = state['step'][state['neuron_idx'] == source]
        count = np.searchsorted(spikes, kc_first[dst[edges]] - 18, side='left')
        label = 'design_ALPN' if source in design else 'non_design_ALPN'
        np.add.at(rows[label + '_strict_pre_first_edge_events'], dst[edges], count)
        np.add.at(rows[label + '_strict_pre_first_positive_V'], dst[edges], count * np.maximum(weights[edges], 0))
        np.add.at(rows[label + '_strict_pre_first_negative_abs_V'], dst[edges], count * np.maximum(-weights[edges], 0))
    return rows


def audit_trial(row, population):
    start = time.perf_counter()
    out = EXP / 'trials' / (row['case'] + '-r' + str(row['repeat']))
    out.mkdir(exist_ok=False)
    source, case = ROOT / row['source'], AG / 'cases' / row['case']
    spec = read(case / 'spec.json')
    graph_name = 'fixtures' if row['kind'] == 'fixture' else 'full'
    core, incoming, raw, codes, targets, positions, graph_audit = graph_view(graph_name)
    with np.load(source / 'state.npz', allow_pickle=False) as data:
        state = {k: data[k].copy() for k in data.files}
    assert set(state) == {'step', 'neuron_idx', 'counts', 'v', 'g', 'ref'}
    with np.load(case / 'requests.npz', allow_pickle=False) as data:
        req = {k: data[k].copy() for k in data.files}
    rt = pd.read_csv(source / 'requests.csv')
    assert np.array_equal(req['tick'], rt.tick) and np.array_equal(req['neuron_idx'], rt.neuron_idx)
    assert (rt.attempted == 1).all() and rt.accepted.isin([0, 1]).all()
    assert np.array_equal(rt.rejected, 1 - rt.accepted)
    if row['kind'] == 'calibration':
        with np.load(case / 'pn-drive.npz', allow_pickle=False) as data:
            design = np.sort(data['targets'].copy())
            assert design.dtype == np.int32 and data['rates_hz'].dtype == np.float32
            assert len(design) == spec['target_count'] and (data['rates_hz'] > 0).all()
    else:
        design = np.intersect1d(req['neuron_idx'], targets)
    assert np.isin(design, targets).all() and len(np.unique(design)) == len(design)
    expected_mask = np.isin(state['neuron_idx'], targets)
    expected = {k: state[k][expected_mask] for k in ('step', 'neuron_idx')}
    initial_v = np.full(len(targets), spec['model']['resting_potential_mv'], np.float64) * np.float64(.001)
    initial_g = np.zeros(len(targets), np.float64)
    if (case / 'arrays.npz').exists():
        with np.load(case / 'arrays.npz', allow_pickle=False) as data:
            initial_v = data['initial_v_mv'][positions] * np.float64(.001)
            initial_g = data['initial_g_mv'][positions] * np.float64(.001)
    histories = {}
    if row['kind'] == 'fixture':
        for key, dtype in (('v', 'float64'), ('g', 'float64'), ('ref', 'int32')):
            histories[key] = np.load(source / f'{key}-history.{dtype}.npy', mmap_mode='r')[:, positions]
            histories['pre_' + key] = np.load(source / f'preAPL-{key}.{dtype}.npy', mmap_mode='r')[:, positions]
    sink = EventSink(out / 'accepted-prefirst-edge-events.parquet')
    observer = PNObserver(targets, design, incoming, codes, GROUPS, read(COEFF), spec['steps'],
                          max_first_event_rows=1000000, event_sink=sink)
    request_lookup = np.full(166700, -1, np.int32)
    request_lookup[targets] = np.arange(len(targets), dtype=np.int32)
    def callback(tick, stage):
        for key in ('v', 'g', 'ref'):
            if histories:
                assert bit_equal(stage[key], histories[key][tick]), 'Fixture post state mismatch'
                assert bit_equal(stage[key + '_pre_apl'], histories['pre_' + key][tick]), 'Fixture pre state mismatch'
        indices = np.flatnonzero((req['tick'] == tick) & np.isin(req['neuron_idx'], targets))
        if len(indices):
            local = request_lookup[req['neuron_idx'][indices]]
            assert np.array_equal(stage['request_accepted'][local], rt.accepted.to_numpy()[indices].astype(bool))
        observer(tick, stage)
    try:
        result = core.replay(initial_v, initial_g, read(COEFF), spec['steps'], targets, req,
            dict(step=state['step'], neuron_idx=state['neuron_idx']), incoming,
            np.zeros((spec['steps'], len(targets)), np.float32), callback)
    finally:
        sink.close()
    exact = {k: bit_equal(result[k], state[k][positions]) for k in ('v', 'g', 'ref', 'counts')}
    exact.update({k: bit_equal(result[k], expected[k]) for k in ('step', 'neuron_idx')})
    assert all(exact.values()), 'PN replay diverged: ' + str(exact)
    data, observer_report = observer.finish(result)
    assert sink.rows == observer_report['first_arrival_record_count']
    table = pd.DataFrame(data)
    table.insert(1, 'body_id', population.iloc[targets].body_id.to_numpy())
    table.insert(2, 'annotation_type', population.iloc[targets].annotation_type.to_numpy())
    table.to_parquet(out / 'per-pn.parquet', index=False)
    edges = pd.DataFrame(observer.edge_summary())
    if len(edges):
        edge_positions = np.searchsorted(incoming.edge_id, edges.edge_id.to_numpy())
        assert np.array_equal(incoming.edge_id[edge_positions], edges.edge_id.to_numpy())
        edges['raw_contact_count'] = raw[edge_positions]
        for prefix in ('source', 'target'):
            global_ids = edges[prefix + '_neuron_idx'].to_numpy()
            edges[prefix + '_body_id'] = population.iloc[global_ids].body_id.to_numpy()
            edges[prefix + '_type'] = population.iloc[global_ids].annotation_type.to_numpy()
    edges.to_parquet(out / 'prefirst-edge-summary.parquet', index=False)
    with (out / 'replay-final.npz').open('xb') as f:
        np.savez_compressed(f, **result)
    kc = np.load(AG / 'apl-bundle/kc_globals.npy')
    pn = np.load(AG / 'apl-bundle/pn_globals.npy')
    first_kc = first_ticks(state['step'], state['neuron_idx'], kc)
    first_pn = first_ticks(state['step'], state['neuron_idx'], pn)
    first_nondesign = min_observed(first_pn[~np.isin(pn, design)])
    first_anykc = min_observed(first_kc)
    component = read(source / 'apl-component.json')
    env = component['environment']
    assert env['cuda_matmul_allow_tf32'] is False and env['float32_matmul_precision'] == 'highest'
    d32 = np.array([int(component['decay_f32_bits'], 16)], np.uint32).view(np.float32)[0]
    q32 = np.array([int(component['gain_f32_bits'], 16)], np.uint32).view(np.float32)[0]
    apl = decompose_apl_sources(np.load(AG / 'apl-bundle/input_matrix.npy', mmap_mode='r'),
        np.load(AG / 'apl-bundle/output_matrix.npy', mmap_mode='r'), pn, kc, design,
        state['step'], state['neuron_idx'], np.load(source / 'apl-drive.float32.npy', mmap_mode='r'),
        np.load(source / 'apl-state.float32.npy', mmap_mode='r'),
        np.load(source / 'apl-inhibition.float32.npy', mmap_mode='r'),
        decay32=d32, gain32=q32, standard_fp32=True, first_kc_tick=first_kc)
    assert apl['summary']['pass_all']
    save(out / 'apl-source-summary.json', apl['summary'])
    pd.DataFrame(apl['tick_rows']).to_parquet(out / 'apl-per-tick.parquet', index=False)
    per_kc = pd.DataFrame(apl['per_kc'])
    per_kc.insert(1, 'body_id', population.iloc[kc].body_id.to_numpy())
    per_kc['KC_present_in_source_graph'] = np.isin(kc, np.load(spec['graph_dir'] + '/nodes.npy'))
    per_kc['first_non_design_PN_tick'] = first_nondesign
    per_kc['first_non_design_relation'] = [temporal_relation(first_nondesign, int(t)) for t in first_kc]
    for key, value in pn_kc_before_first(graph_name, state, first_kc, design).items():
        if key not in ('neuron_idx', 'first_spike_tick'):
            per_kc[key] = value
    per_kc.to_parquet(out / 'per-kc-apl-sources.parquet', index=False)
    report = dict(case=row['case'], repeat=row['repeat'], kind=row['kind'], pass_all=True, exact=exact,
        every_tick_fixture_pre_post_exact=bool(histories), CAL_per_tick_independent_state_available=False,
        steps=spec['steps'], duration_ms=spec['duration_ms'], mode=spec.get('mode'), seed=spec.get('seed'),
        odor_id=spec.get('odor_id'), graph=graph_audit, observer=observer_report,
        first_non_design_PN_tick=first_nondesign, first_any_KC_tick=first_anykc,
        first_non_design_vs_first_KC=temporal_relation(first_nondesign, first_anykc),
        fixture_design_definition='Unique requested PN IDs for numerical fixture only; not odor design' if histories else None,
        APL_source_bound_pass=apl['summary']['pass_all'], candidate_generated=False, new_brain_simulation=False,
        CPU_audit_wall_seconds=time.perf_counter() - start,
        input_sha256={f: digest(source / f) for f in ('state.npz', 'requests.csv', 'summary.json',
            'apl-component.json', 'apl-drive.float32.npy', 'apl-state.float32.npy', 'apl-inhibition.float32.npy')})
    save(out / 'summary.json', report)
    save(out / 'output-hashes.json', [dict(path=p.name, sha256=digest(p), bytes=p.stat().st_size)
         for p in sorted(out.iterdir()) if p.is_file()])
    return report


def main():
    check()
    population = pd.read_parquet(EXP / 'protocol/neuron-source-groups.parquet')
    inventory = read(EXP / 'protocol/inventory.json')
    receipts = []
    for row in inventory:
        try:
            report = audit_trial(row, population)
            receipts.append(dict(case=row['case'], repeat=row['repeat'], pass_all=report['pass_all']))
            assert sum(p.stat().st_size for p in EXP.rglob('*') if p.is_file()) < 2 * 1024**3
            print(row['case'], row['repeat'], 'EXACT PN / APL BOUND PASS',
                  round(report['CPU_audit_wall_seconds'], 3), 'CPU sec', flush=True)
        except Exception as error:
            save(EXP / 'audit/failure.json', dict(case=row, error=type(error).__name__, message=str(error),
                 traceback=traceback.format_exc(), completed=receipts, new_brain_simulation=False))
            raise
    repeats = []
    for row in inventory:
        if row['repeat'] != 1:
            continue
        a = EXP / 'trials' / (row['case'] + '-r0')
        b = EXP / 'trials' / (row['case'] + '-r1')
        checks = {}
        for f in ('per-pn.parquet', 'prefirst-edge-summary.parquet', 'apl-per-tick.parquet', 'per-kc-apl-sources.parquet'):
            checks[f] = pd.read_parquet(a / f).equals(pd.read_parquet(b / f))
        checks['antecedent_events'] = digest(a / 'accepted-prefirst-edge-events.parquet') == digest(b / 'accepted-prefirst-edge-events.parquet')
        checks['apl-summary'] = read(a / 'apl-source-summary.json') == read(b / 'apl-source-summary.json')
        with np.load(a / 'replay-final.npz') as aa, np.load(b / 'replay-final.npz') as bb:
            checks['replay'] = set(aa.files) == set(bb.files) and all(bit_equal(aa[k], bb[k]) for k in aa.files)
        assert all(checks.values()), 'Repeated accounting differs'
        repeats.append(dict(case=row['case'], pass_all=True, checks=checks))
    assert len(receipts) == 80 and len(repeats) == 12
    save(EXP / 'audit/completed.json', dict(pass_all=True, trials=receipts, repeat_checks=repeats,
         meaning='Conditional PN exactness and bounded APL attribution only; not a learning or causal gate'))
    print('PASS 80 conditional audits and 12 repeats', flush=True)


if __name__ == '__main__':
    main()