"""Whole-release leaky integrate-and-fire experiments, without a behavior policy.

Equations/parameters follow Shiu et al. 2024's simple model. Applying them to
MaleCNS/BANC is a new modeling experiment, not a reproduction of their result.
All retained edges remain instantiated, including ambiguous transmitter edges.
"""
from dataclasses import dataclass
from pathlib import Path
import json
import hashlib
import time

import numpy as np

ROOT = Path(__file__).resolve().parents[1]


def transmitter_signs(transmitters, ambiguous_sign=1):
    """Declared coarse fast-transmission assumption; never deletes unknown edges.

    ACh +; GABA, glutamate, histamine -. A co-transmitter combination with only
    one fast sign uses that sign; conflicting, missing and modulator-only cells
    use the explicit sensitivity parameter. This is NOT receptor physiology.
    """
    if ambiguous_sign not in (-1, 1):
        raise ValueError("Ambiguous edges must remain active with sign +1 or -1.")
    signs, uncertain = [], []
    for value in transmitters:
        tokens = set(str(value).lower().split(','))
        fast = ({1} if 'acetylcholine' in tokens else set()) | ({-1} if tokens & {'gaba', 'glutamate', 'histamine'} else set())
        ambiguous = len(fast) != 1
        signs.append(ambiguous_sign if ambiguous else next(iter(fast)))
        uncertain.append(ambiguous)
    return np.asarray(signs, dtype=np.int8), np.asarray(uncertain, dtype=bool)


@dataclass
class WholeConnectome:
    dataset: str
    animal_id: str

    def __post_init__(self):
        import brian2 as b2
        import pyarrow as pa
        import pyarrow.feather as feather
        import pyarrow.ipc as ipc
        self.b2 = b2
        path = ROOT / 'connectome_data' / self.dataset / 'normalized'
        self.nodes = feather.read_table(path / 'neurons.feather').to_pandas()
        self.import_report = json.loads((path / 'report.json').read_text())
        self.ids = self.nodes.source_id.to_numpy(dtype=np.int64)
        edges = ipc.open_file(pa.memory_map(str(path / 'edges.arrow'), 'r')).read_all()
        pre, post, count = [edges.column(c).to_numpy() for c in ('pre_index', 'post_index', 'synapse_count')]
        assert len(pre) == self.import_report['graph']['retained_edge_rows']
        assert int(count.sum(dtype=np.uint64)) == self.import_report['graph']['retained_synaptic_contacts']
        b2.prefs.codegen.target = 'cython'
        b2.prefs.codegen.runtime.cython.cache_dir = str(ROOT / '.cache-neural')
        self.clock = b2.Clock(dt=.1 * b2.ms)
        # Like the published model, g is a voltage-valued synaptic drive.
        self.neurons = b2.NeuronGroup(len(self.ids), '''
            dv/dt = (-52*mV - v + g)/(20*ms) : volt (unless refractory)
            dg/dt = -g/(5*ms) : volt (unless refractory)
            rfc : second
            ''', threshold='v > -45*mV', reset='v = -52*mV; g = 0*mV',
            refractory='rfc', method='linear', clock=self.clock)
        self.neurons.v = -52 * b2.mV
        self.neurons.g = 0 * b2.mV
        self.neurons.rfc = 2.2 * b2.ms
        self.synapses = b2.Synapses(self.neurons, self.neurons, 'w : volt',
                                   on_pre='g_post += w', delay=1.8*b2.ms, clock=self.clock)
        self.synapses.connect(i=pre.astype(np.int32), j=post.astype(np.int32))
        self.pre, self.count = pre, count
        self.monitor = b2.SpikeMonitor(self.neurons)
        # One independently sampled Poisson channel for each selected input.
        self.stimulus = b2.PoissonGroup(len(self.ids), rates=0*b2.Hz, clock=self.clock)
        self.stimulus_monitor = b2.SpikeMonitor(self.stimulus)
        self.input_synapses = b2.Synapses(self.stimulus, self.neurons,
                                        on_pre='v_post += 68.75*mV', clock=self.clock)
        self.input_synapses.connect(j='i')
        self.network = b2.Network(self.neurons, self.synapses, self.monitor, self.stimulus, self.stimulus_monitor, self.input_synapses)
        self.set_signs(1)
        self.network.store('rest')

    def set_signs(self, ambiguous_sign):
        signs, uncertain = transmitter_signs(self.nodes.neurotransmitter, ambiguous_sign)
        self.synapses.w = self.count * signs[self.pre] * .275 * self.b2.mV
        self.uncertain = uncertain

    def indices(self, cell_types):
        return np.flatnonzero(self.nodes.cell_type.isin(cell_types).to_numpy())

    def run(self, condition, cell_types=(), duration_ms=1000, rate_hz=100,
            seed=41027, disconnect=False, ambiguous_sign=1):
        b2 = self.b2
        if duration_ms <= 0 or rate_hz < 0:
            raise ValueError('Invalid duration or stimulation rate')
        selected = self.indices(cell_types)
        if cell_types and not len(selected):
            raise ValueError(f'No annotated input cells for {cell_types}')
        self.network.restore('rest')
        self.set_signs(ambiguous_sign)
        b2.seed(seed)
        self.stimulus.rates = 0 * b2.Hz
        self.stimulus.rates[selected] = rate_hz * b2.Hz
        self.neurons.rfc[selected] = 0 * b2.ms
        cut_edges = 0
        if disconnect:
            # Explicit causal intervention, never a graph-import filter.
            mask = np.isin(self.pre, selected)
            cut_edges = int(mask.sum())
            self.synapses.w[np.flatnonzero(mask)] = 0 * b2.mV
        began = time.monotonic()
        self.network.run(duration_ms*b2.ms, namespace={})
        indices = np.asarray(self.monitor.i)
        times = np.asarray(self.monitor.t/b2.ms)
        counts = np.asarray(self.monitor.count)
        input_indices=np.asarray(self.stimulus_monitor.i,dtype='<i8')
        input_times=np.asarray(self.stimulus_monitor.t/b2.ms,dtype='<f8')
        input_digest=hashlib.sha256(input_indices.tobytes()+input_times.tobytes()).hexdigest()
        targets = self.indices(['DNa01','DNa02','MN9','MN10','MN11D','MN11V','MN12D','CEM'])
        bins = np.arange(0, duration_ms + 10, 10)
        return {
            'condition': condition, 'duration_ms': duration_ms, 'seed': seed,
            'input_cell_types': list(cell_types), 'input_ids': [str(self.ids[i]) for i in selected],
            'input_rate_hz': rate_hz if len(selected) else 0,
            'ambiguous_sign': ambiguous_sign, 'silenced_outgoing_edges': cut_edges,
            'total_spikes': int(counts.sum()), 'active_neurons': int(np.count_nonzero(counts)),
            'input_spikes': int(counts[selected].sum()),
            'external_input_spikes': len(input_indices),
            'external_input_events_sha256': input_digest,
            'motor_spikes': int(counts[self.nodes.superclass.str.contains('motor', na=False)].sum()),
            'wall_seconds': round(time.monotonic()-began, 3),
            'population_bins_10ms': np.histogram(times, bins)[0].tolist(),
            'readouts': [{'id': str(self.ids[i]), 'cell_type': self.nodes.cell_type.iloc[i],
                          'spikes': int(counts[i]), 'spike_times_ms': times[indices==i].round(3).tolist()}
                         for i in targets],
            'active_by_superclass': {str(k): int(v) for k,v in self.nodes.assign(spikes=counts).groupby('superclass').spikes.sum().items()},
        }


def run_suite(dataset):
    model = WholeConnectome(dataset, 'male-1' if dataset == 'malecns_v1' else 'female-1')
    out = ROOT / 'outputs' / 'whole-connectome' / dataset
    out.mkdir(parents=True, exist_ok=True)
    results = []
    experiments = [
        ('no_input', (), False, 1),
        ('ORN_DM1_100hz', ('ORN_DM1',), False, 1),
        ('ORN_DM1_outgoing_disconnected', ('ORN_DM1',), True, 1),
        ('LB3c_100hz', ('LB3c',), False, 1),
        ('LB3c_outgoing_disconnected', ('LB3c',), True, 1),
        ('LB3c_ambiguous_negative', ('LB3c',), False, -1),
    ]
    report = {'schema_version': 1, 'dataset': dataset, 'animal_id': model.animal_id,
              'neuron_count': len(model.ids), 'edge_count': len(model.pre),
              'synaptic_contacts': int(model.count.sum(dtype=np.uint64)),
              'uncertain_fast_sign_neurons': int(model.uncertain.sum()),
              'equation_source': 'https://www.nature.com/articles/s41586-024-07763-9',
              'dt_ms': .1, 'delay_ms': 1.8, 'synaptic_gain_mv': .275,
              'biologically_validated': False, 'embodied': False,
              'interpretation': 'New whole-release neural perturbation experiments. LB3c is sugar annotated in BANC; male identity is a cell-type homology inference. ORN_DM1 stimulation is not calibrated banana chemistry. No movement or intake is inferred from motor spikes alone.',
              'sign_assumptions': 'ACh excitatory; GABA/glutamate/histamine inhibitory; unique fast sign takes precedence over co-modulators; all ambiguous/modulator-only/missing signs are +1 with a -1 sensitivity run. No receptor-dependent or hormonal state.',
              'source_hashes': model.import_report['source_hashes'], 'results': results}
    for name, types, disconnect, sign in experiments:
        print(f'{dataset}: {name}', flush=True)
        result = model.run(name, types, disconnect=disconnect, ambiguous_sign=sign)
        results.append(result)
        (out/'report.json').write_text(json.dumps(report, indent=2)+'\n')
        print(json.dumps({k:result[k] for k in ['condition','total_spikes','active_neurons','motor_spikes','wall_seconds']}), flush=True)
    return report


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('dataset', choices=['malecns_v1','banc_v888'])
    run_suite(parser.parse_args().dataset)
