"""All predeclared pairs and all trial outputs; never choose a winning odor."""
from itertools import combinations
import numpy as np
import pandas as pd
from context import EXP, AK, CFG, read, save as _save, digest, bit_equal, write_npz as _write_npz
from credit import pair_identity


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


def metrics(values):
    values = np.asarray(values, np.float64)
    if not len(values):
        return dict(min=None, mean=None, max=None, undefined_reason='No defined denominator observations')
    return dict(min=float(values.min()), mean=float(values.mean()), max=float(values.max()))


def primary_rows(reg):
    rows = [r for r in reg['planned'] if r['repeat'] == 0]
    assert len(rows) == CFG['primary_trials'] and len({r['name'] for r in rows}) == len(rows)
    return rows


def run(verify_only=False):
    def save(path, value):
        if verify_only:
            assert read(path) == value, 'Summary JSON differs from recomputed evidence: ' + str(path)
        else:
            _save(path, value)

    def write_npz(path, **value):
        if verify_only:
            old = load_arrays(path)
            assert set(old) == set(value) and all(bit_equal(old[k], value[k]) for k in old)
        else:
            _write_npz(path, **value)

    reg = read(EXP / 'protocol/preregistered.json')
    assert read(EXP / 'audit/replays-complete.json')['pass_all']
    data = load_arrays(EXP / 'protocol/incoming.npz')
    for row in reg['planned']:
        folder = EXP / 'trials' / (row['name'] + '-r' + str(row['repeat']))
        hashes = read(folder / 'output-hashes.json')
        assert {r['path'] for r in hashes} == {p.name for p in folder.iterdir() if p.name != 'output-hashes.json'}
        for entry in hashes:
            assert digest(folder / entry['path']) == entry['sha256']
        original = load_arrays(AK / 'gpu' / row['cohort'] / (row['name'] + '-r' + str(row['repeat'])) / 'state.npz')
        base = load_arrays(folder / 'baseline-replay.npz')
        selected = np.isin(original['neuron_idx'], data['targets'])
        assert all(bit_equal(base[k], original[k][data['targets']]) for k in ('v', 'g', 'ref', 'counts'))
        assert all(bit_equal(base[k], original[k][selected]) for k in ('step', 'neuron_idx'))
    candidate = pd.read_parquet(EXP / 'protocol/candidate-static.parquet')
    target_meta = pd.read_parquet(EXP / 'protocol/mbon-metadata.parquet')
    primary = primary_rows(reg)
    repeats = []
    for row in reg['planned']:
        if row['repeat'] != 1:
            continue
        paths = [EXP / 'trials' / (row['name'] + '-r' + str(i)) for i in (0, 1)]
        checks = {}
        for filename in ('baseline-replay.npz', 'baseline-target-accounting.npz', 'baseline-edge-accounting.npz',
                'conditional25-replay.npz', 'conditional25-target-accounting.npz', 'conditional25-edge-accounting.npz', 'credit.npz'):
            a, b = (load_arrays(p / filename) for p in paths)
            checks[filename] = set(a) == set(b) and all(bit_equal(a[k], b[k]) for k in a)
        for filename in ('baseline-accounting.json', 'conditional25-accounting.json', 'credit.json', 'sensitivity.json'):
            checks[filename] = read(paths[0] / filename) == read(paths[1] / filename)
        assert all(checks.values()), 'Repeat accounting differs: ' + row['name']
        repeats.append(dict(name=row['name'], checks=checks, pass_all=True))
    assert len(repeats) == CFG['technical_repeats']
    save(EXP / 'audit/repeat-checks.json', dict(pass_all=True, repeats=repeats))
    summaries, target_rows, trial_metrics = [], [], []
    for row in primary:
        folder = EXP / 'trials' / (row['name'] + '-r0')
        report = read(folder / 'summary.json')
        assert report['pass_all'] and all(report['baseline_exact'].values())
        summaries.append(report)
        targets = load_arrays(folder / 'baseline-target-accounting.npz')
        change = load_arrays(folder / 'conditional25-target-accounting.npz')
        edges = load_arrays(folder / 'baseline-edge-accounting.npz')
        exc = targets['accepted_positive_V'].sum(axis=1)
        plastic_exc = targets['accepted_positive_V'][:, 0]
        area_exc = targets['g_sign_integral_V_ms'][:, :, 0].sum(axis=1)
        area_plastic = targets['g_sign_integral_V_ms'][:, 0, 0]
        for i, meta in enumerate(target_meta.itertuples()):
            entry = dict(name=row['name'], cohort=row['cohort'], mode=row['mode'], odor_id=row['odor_id'], seed=row['seed'],
                neuron_idx=int(meta.neuron_idx), body_id=int(meta.body_id), annotation_type=meta.annotation_type,
                compartment_verbatim=meta.mb_compartment_verbatim,
                candidate_target=bool(np.any(candidate.target_neuron_idx.to_numpy() == meta.neuron_idx)),
                accepted_total_positive_V=float(exc[i]), accepted_plastic_positive_V=float(plastic_exc[i]),
                plastic_accepted_positive_fraction=None if exc[i] == 0 else float(plastic_exc[i] / exc[i]),
                positive_source_g_area_V_ms=float(area_exc[i]), plastic_positive_g_area_V_ms=float(area_plastic[i]),
                plastic_positive_g_area_fraction=None if area_exc[i] == 0 else float(area_plastic[i] / area_exc[i]),
                baseline_spikes=int(targets['spike_count'][i]), conditional25_spikes=int(change['spike_count'][i]),
                baseline_first_tick=int(targets['first_spike_tick'][i]), conditional25_first_tick=int(change['first_spike_tick'][i]),
                max_abs_V_delta_SI=float(np.max(np.abs(targets['history_v'][:, i] - change['history_v'][:, i]))),
                max_abs_g_delta_SI=float(np.max(np.abs(targets['history_g'][:, i] - change['history_g'][:, i]))))
            target_rows.append(entry)
        pick = data['is_candidate']
        summary = dict(name=row['name'], cohort=row['cohort'], mode=row['mode'], odor_id=row['odor_id'], seed=row['seed'],
            plastic_due_events=int(edges['due_count'][pick].sum()), plastic_accepted_events=int(edges['accepted_count'][pick].sum()),
            plastic_rejected_events=int(edges['rejected_count'][pick].sum()),
            plastic_pending_events=int(edges['pending_after_window_count'][pick].sum()),
            plastic_edges_with_accepted_events=int(np.count_nonzero(edges['accepted_count'][pick])),
            accepted_positive_V=float(exc.sum()), plastic_accepted_positive_V=float(plastic_exc.sum()),
            plastic_accepted_positive_fraction=None if exc.sum() == 0 else float(plastic_exc.sum() / exc.sum()),
            plastic_positive_g_area_fraction=None if area_exc.sum() == 0 else float(area_plastic.sum() / area_exc.sum()),
            **report['sensitivity'])
        trial_metrics.append(summary)
    target_table = pd.DataFrame(target_rows)
    if verify_only:
        assert target_table.equals(pd.read_parquet(EXP / 'audit/per-MBON-primary.parquet'))
    else:
        target_table.to_parquet(EXP / 'audit/per-MBON-primary.parquet', index=False)
    save(EXP / 'audit/primary-trial-metrics.json', trial_metrics)
    pair_reports = []
    for cohort in sorted({r['cohort'] for r in primary}):
        for mode in sorted({r['mode'] for r in primary}):
            rows = [r for r in primary if r['cohort'] == cohort and r['mode'] == mode]
            odors = sorted({r['odor_index'] for r in rows})
            assert len(rows) == 30 and len(odors) == 6
            traces = {}; whole = {}; ids = {}
            for odor in odors:
                series = sorted((r for r in rows if r['odor_index'] == odor), key=lambda r: r['seed'])
                assert len(series) == 5
                ids[odor] = series[0]['odor_id']
                credit = [load_arrays(EXP / 'trials' / (r['name'] + '-r0') / 'credit.npz') for r in series]
                traces[odor] = np.stack([c['trace_gate'] for c in credit])
                whole[odor] = np.stack([(c['spike_count_total'] > 0).astype(np.float32) for c in credit])
            for a, b in combinations(odors, 2):
                name = f'{cohort}-{mode}-{a:02d}-vs-{b:02d}'
                result = dict(cohort=cohort, mode=mode, odor_a=ids[a], odor_b=ids[b], odor_index_a=a, odor_index_b=b)
                for definition, series in (('virtual_gate900', traces), ('whole_window_auxiliary', whole)):
                    identities = pair_identity(series[a], series[b], data['edge_source_local'])
                    save_arrays = {k: v for k, v in identities.items() if isinstance(v, np.ndarray)}
                    write_npz(EXP / 'pairs' / (name + '-' + definition + '.npz'), **save_arrays)
                    summary = identities['summary']
                    # Historic local existence criterion only, no new overall pass rate.
                    summary['local_specific_KC_and_edges_exist_both_sides'] = all(
                        summary[unit]['category_counts'][side] >= 1 for unit in ('KC', 'edges') for side in ('A_specific', 'B_specific'))
                    summary['compartments'] = {}
                    for comp, group in candidate.groupby('target_compartment_verbatim', sort=True):
                        sub = pair_identity(series[a], series[b], data['edge_source_local'][group.index.to_numpy()])
                        summary['compartments'][str(comp)] = sub['summary']['edges']
                    result[definition] = summary
                pair_reports.append(result)
    assert len(pair_reports) == 60
    save(EXP / 'audit/pair-identities.json', pair_reports)
    blocks = {}
    for cohort in sorted({r['cohort'] for r in primary}):
        for mode in sorted({r['mode'] for r in primary}):
            trials = [r for r in trial_metrics if r['cohort'] == cohort and r['mode'] == mode]
            pairs = [p for p in pair_reports if p['cohort'] == cohort and p['mode'] == mode]
            blocks[cohort + '/' + mode] = dict(trials=len(trials), pairs=len(pairs),
                pairs_with_both_specific_KC_and_edge_sets=sum(p['virtual_gate900']['local_specific_KC_and_edges_exist_both_sides'] for p in pairs),
                KC_empty_union_pairs=sum(p['virtual_gate900']['KC']['empty_union'] for p in pairs),
                edge_empty_union_pairs=sum(p['virtual_gate900']['edges']['empty_union'] for p in pairs),
                KC_active_set_Jaccard=metrics([p['virtual_gate900']['KC']['active_set_jaccard'] for p in pairs if not p['virtual_gate900']['KC']['empty_union']]),
                candidate_edge_active_set_Jaccard=metrics([p['virtual_gate900']['edges']['active_set_jaccard'] for p in pairs if not p['virtual_gate900']['edges']['empty_union']]),
                plastic_accepted_events=metrics([r['plastic_accepted_events'] for r in trials]),
                plastic_fraction_of_all_MBON_accepted_positive=metrics([r['plastic_accepted_positive_fraction'] for r in trials if r['plastic_accepted_positive_fraction'] is not None]),
                conditional25_event_changed_trials=sum(r['event_sequence_changed'] for r in trials),
                conditional25_count_or_first_latency_changed_trials=sum(r['functional_count_or_first_latency_change'] for r in trials),
                max_abs_V_difference_SI=max(r['max_abs_V_difference_SI'] for r in trials),
                max_abs_g_difference_SI=max(r['max_abs_g_difference_SI'] for r in trials))
    outcome = dict(arithmetic_and_coverage_pass=True, primary_trials=120, exact_repeats=24, pairs=60, blocks=blocks,
        candidate_generated=False, learning_achieved=False, conditioning_executed=False, scientific_overall_PASS_threshold_defined=False,
        maximum_observed_g_lane_roundoff_residual_SI=max(read(EXP / 'trials' / (r['name'] + '-r0') / 'baseline-accounting.json')['max_g_lane_residual_V'] for r in primary),
        interpretation='Read-only credit delivery and fixed conditional sensitivity; not closed-loop memory or dopamine measurement')
    save(EXP / 'audit/scientific-summary.json', outcome)
    return outcome


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