"""Descriptive post-audit metrics from sealed outputs, without new trajectories."""
from pathlib import Path
import json
import sys
import numpy as np
import pandas as pd

ROOT = Path('/data/digital-drosophila')
EXP = ROOT / 'experiments/phase4ah-kc-drive-audit'
MODEL = ROOT / 'models/male-cns-lif/kc-drive-audit-v1'
sys.path.insert(0, str(MODEL))
from context import read, save, digest
from populations import GROUPS


def ratio(a, b):
    return None if b == 0 else float(a / b)


def main():
    final = read(EXP / 'audit/final-status.json')
    assert final['status'] == 'PASS-CONDITIONAL-KC-DRIVE-AUDIT'
    manifest = {r['path']: r for r in read(EXP / 'manifest.json')['files']}
    grouped = {}
    inputs = []
    for item in read(EXP / 'protocol/inventory.json'):
        if item['kind'] != 'calibration' or item['repeat'] != 0:
            continue
        name = item['case'] + '-r0'
        folder = EXP / 'trials' / name
        for filename in ('summary.json', 'per-kc.parquet', 'per-pn.parquet'):
            p = folder / filename
            rel = str(p.relative_to(EXP))
            assert digest(p) == manifest[rel]['sha256']
            inputs.append(dict(path=rel, sha256=manifest[rel]['sha256']))
        summary = read(folder / 'summary.json')
        kc = pd.read_parquet(folder / 'per-kc.parquet')
        pn = pd.read_parquet(folder / 'per-pn.parquet')
        grouped.setdefault(summary['mode'], []).append((summary, kc, pn))
    output = {}
    for mode, rows in grouped.items():
        assert len(rows) == 30
        allkc = pd.concat([x[1] for x in rows], ignore_index=True)
        assert len(allkc) == 30 * 4064
        strata = {}
        for label, data in (('all', allkc), ('ever_active', allkc[allkc.spike_count > 0]),
                            ('never_active', allkc[allkc.spike_count == 0])):
            exc = data[[g + '_accepted_exc_SI' for g in GROUPS]].sum(axis=1)
            inh = data[[g + '_accepted_inh_abs_SI' for g in GROUPS]].sum(axis=1)
            pnexc = data.ALPN_accepted_exc_SI
            margin = data.best_threshold_margin_mV
            strata[label] = dict(neuron_trials=len(data),
                with_ALPN_exc_neuron_trials=int((pnexc > 0).sum()),
                with_any_exc_neuron_trials=int((exc > 0).sum()),
                with_ALPN_exc_fraction=ratio(int((pnexc > 0).sum()), len(data)),
                with_any_exc_fraction=ratio(int((exc > 0).sum()), len(data)),
                accepted_exc_increment_sum_V=float(exc.sum()),
                accepted_inh_increment_abs_sum_V=float(inh.sum()),
                ALPN_exc_increment_sum_V=float(pnexc.sum()),
                APL_write_increment_sum_V=float(data.APL_write_SI.sum()),
                APL_write_ref_sum_V=float(data.APL_write_ref_SI.sum()),
                APL_pre_first_spike_increment_sum_V=float(data.APL_pre_first_spike_SI.sum()),
                APL_write_to_exc_increment_ratio=ratio(data.APL_write_SI.sum(), exc.sum()),
                refractory_target_ticks=int(data.refractory_ticks.sum()),
                negative_g_target_tick_fraction=ratio(data.negative_g_ticks.sum(), len(data) * 1000),
                best_threshold_margin_mV_quantiles={str(q): float(margin.quantile(q)) for q in (0, .25, .5, .75, 1)},
                source_g_area_V_ms={g: float(data['source_g_area_' + g + '_V_ms'].sum())
                    for g in [n + '_' + sign for n in GROUPS for sign in ('exc', 'inh')] + ['graded_APL']},
                source_area_max_abs_residual_V_ms=float(data.source_g_area_residual_V_ms.abs().max()))
        accepted = sum(x[0]['input']['accepted'] for x in rows)
        requested = sum(x[0]['input']['requests'] for x in rows)
        source = final['modes'][mode]['source_groups']
        source_exc = {g: source[g]['accepted_exc_SI'] for g in GROUPS}
        output[mode] = dict(trials=30, input_requests=requested, input_accepted=accepted,
            input_acceptance_fraction=ratio(accepted, requested),
            PN_actual_spikes=sum(x[0]['input']['actual_spikes_by_source_group']['ALPN'] for x in rows),
            PN_design_targets_distinct=sorted({int(i) for _, _, pn in rows for i in pn.loc[pn.design_target, 'neuron_idx']}),
            PN_ever_active_distinct=sorted({int(i) for _, _, pn in rows for i in pn.loc[pn.actual_spikes > 0, 'neuron_idx']}),
            source_accepted_exc_increment_sum_V=source_exc,
            ALPN_fraction_of_accepted_exc_increment=ratio(source_exc['ALPN'], sum(source_exc.values())),
            ALPN_due_edges=source['ALPN']['due_edges'], ALPN_accepted_edges=source['ALPN']['accepted_edges'],
            ALPN_rejected_edges=source['ALPN']['rejected_edges'],
            ALPN_rejection_fraction=ratio(source['ALPN']['rejected_edges'], source['ALPN']['due_edges']),
            strata=strata)
    result = dict(scope='Post-outcome descriptive conditional audit, no causal intervention, no learning', modes=output,
        target_tick_and_neuron_trial_counts_not_unique_cells=True,
        ratios_are_cumulative_increment_bookkeeping_not_spike_causality=True,
        inputs=inputs, source_script_sha256=digest(Path(__file__)),
        sealed_model_manifest_sha256=digest(MODEL / 'manifest.json'),
        sealed_experiment_manifest_sha256=digest(EXP / 'manifest.json'))
    destination = ROOT / 'registry/phase4ah-kc-drive-audit/interpretation-metrics.json'
    save(destination, result)
    with destination.with_suffix('.sha256').open('x') as f:
        f.write(digest(destination) + '  ' + destination.name + '\n')
    print(json.dumps({mode: {k: v for k, v in data.items() if k not in ('PN_design_targets_distinct', 'PN_ever_active_distinct')}
                      for mode, data in output.items()}, indent=2), flush=True)


if __name__ == '__main__':
    main()