"""Complete paired static-overlay effects; no learned-memory inference."""
import numpy as np
import pandas as pd
from context import EXP, CFG, read, save, digest, bit_equal
from run_suite import verify_completed


def load(path):
    with np.load(path, allow_pickle=False) as data:
        return {k: data[k].copy() for k in data.files}


def first_ticks(state, n=166700):
    out = np.full(n, np.iinfo(np.int32).max, np.int32)
    np.minimum.at(out, state['neuron_idx'], state['step'])
    out[out == np.iinfo(np.int32).max] = -1
    return out


def metrics(base, changed, population):
    a, b = base['counts'][population], changed['counts'][population]
    af, bf = first_ticks(base)[population], first_ticks(changed)[population]
    ae = np.isin(base['neuron_idx'], population); be = np.isin(changed['neuron_idx'], population)
    same = bit_equal(base['step'][ae], changed['step'][be]) and bit_equal(base['neuron_idx'][ae], changed['neuron_idx'][be])
    ax = base['step'][ae].astype(np.int64) * 166700 + base['neuron_idx'][ae]
    bx = changed['step'][be].astype(np.int64) * 166700 + changed['neuron_idx'][be]
    different = np.setxor1d(ax, bx, assume_unique=True)
    return dict(cells=len(population), base_active=int(np.count_nonzero(a)), changed_active=int(np.count_nonzero(b)),
        base_spikes=int(a.sum()), changed_spikes=int(b.sum()),
        recruited_cells=int(np.count_nonzero((a == 0) & (b > 0))), lost_cells=int(np.count_nonzero((a > 0) & (b == 0))),
        count_changed_cells=int(np.count_nonzero(a != b)), first_spike_changed_cells=int(np.count_nonzero(af != bf)),
        count_l1=int(np.abs(a.astype(np.int64) - b.astype(np.int64)).sum()), event_sequence_changed=not same,
        first_event_divergence_ms=None if not len(different) else float(different.min() // 166700 * .1),
        final_v_max_abs_SI=float(np.max(np.abs(base['v'][population] - changed['v'][population]))) if len(population) else None,
        final_g_max_abs_SI=float(np.max(np.abs(base['g'][population] - changed['g'][population]))) if len(population) else None)


def run(verify_only=False):
    def emit(path, value):
        if verify_only:
            assert read(path) == value, str(path)
        else:
            save(path, value)
    reg = read(EXP / 'protocol/preregistered.json')
    for stage in CFG['stages']:
        receipt = read(EXP / 'audit' / (stage + '-complete.json'))
        assert receipt['pass_all'] and receipt['trials'] == 144
        assert len(receipt['completed']) == 144
        for row, entry in zip(reg['planned'], receipt['completed']):
            assert entry['name'] == row['name'] and entry['repeat'] == row['repeat']
            folder = EXP / 'gpu' / stage / (row['name'] + '-r' + str(row['repeat']))
            verify_completed(folder, stage, row, recompute=True)
            assert digest(folder / 'output-hashes.json') == entry['output_hashes_sha256']
            assert digest(folder / 'comparison.json') == entry['comparison_sha256']
    pops = load(EXP / 'protocol/populations.npz')
    rows, cellrows, times, peaks = [], [], [], []
    for row in reg['planned']:
        if row['repeat']:
            continue
        key = row['name'] + '-r0'
        base, changed = [load(EXP / 'gpu' / stage / key / 'state.npz') for stage in CFG['stages']]
        effects = {name: metrics(base, changed, nodes) for name, nodes in pops.items()}
        result = dict(**row, effects=effects, whole_brain=metrics(base, changed, np.arange(166700, dtype=np.int32)))
        emit(EXP / 'summaries' / (row['name'] + '.json'), result); rows.append(result)
        af, bf = first_ticks(base), first_ticks(changed)
        for neuron in pops['MBON']:
            n = int(neuron)
            cellrows.append(dict(case=row['name'], cohort=row['cohort'], mode=row['mode'], seed=row['seed'],
                neuron_idx=n, base_count=int(base['counts'][n]), static25_count=int(changed['counts'][n]),
                base_first_tick=int(af[n]), static25_first_tick=int(bf[n]),
                base_final_v_SI=float(base['v'][n]), static25_final_v_SI=float(changed['v'][n]),
                base_final_g_SI=float(base['g'][n]), static25_final_g_SI=float(changed['g'][n])))
        for stage in CFG['stages']:
            s = read(EXP / 'gpu' / stage / key / 'summary.json')
            times.append(dict(stage=stage, case=row['name'], wall_seconds=s['wall_seconds'],
                simulation_seconds=s['diagnostic_simulation_seconds']))
            peaks.append(s['peak_allocated_vram_bytes'])
    table = pd.DataFrame(cellrows)
    table_path = EXP / 'summaries/per-MBON-static-effects.parquet'
    if verify_only:
        pd.testing.assert_frame_equal(table, pd.read_parquet(table_path), check_exact=True)
    else:
        assert not table_path.exists(); table.to_parquet(table_path, index=False)
    blocks = {}
    for cohort in sorted({r['cohort'] for r in rows}):
        for mode in sorted({r['mode'] for r in rows}):
            block = [r for r in rows if r['cohort'] == cohort and r['mode'] == mode]
            assert len(block) == 30
            blocks[cohort + '/' + mode] = dict(trials=30, populations={name: dict(
                trials_event_change=sum(r['effects'][name]['event_sequence_changed'] for r in block),
                trials_count_change=sum(r['effects'][name]['count_changed_cells'] > 0 for r in block),
                mean_count_l1=float(np.mean([r['effects'][name]['count_l1'] for r in block])),
                mean_recruited_cells=float(np.mean([r['effects'][name]['recruited_cells'] for r in block])),
                mean_lost_cells=float(np.mean([r['effects'][name]['lost_cells'] for r in block]))) for name in pops})
    summary = dict(status='PASS-EXACT-WHOLEBRAIN-PLASTIC-INTERFACE-NOT-LEARNING',
        primary_trials_per_stage=120, technical_repeats_per_stage=24, zero_overlay_exact=True,
        static25_repeat_exact=True, blocks=blocks, peak_allocated_vram_bytes=max(peaks),
        timing=times, per_MBON_sha256=digest(table_path), training=False, learned_candidate=False,
        conditional_AL_not_a_nonzero_arithmetic_oracle=True,
        scope='Valid interface and recording; reported static response is not an associative learning PASS')
    emit(EXP / 'audit/scientific-summary.json', summary)
    return summary


if __name__ == '__main__':
    print(run()['status'], flush=True)