ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.biosimulator

    # Copyright 2017. Allen Institute. All rights reserved
    #
    # Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
    # following conditions are met:
    #
    # 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
    # disclaimer.
    #
    # 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
    # disclaimer in the documentation and/or other materials provided with the distribution.
    #
    # 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
    # products derived from this software without specific prior written permission.
    #
    # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
    # INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    #
    import time
    from six import string_types
    from neuron import h
    from bmtk.simulator.core.simulator import Simulator
    from bmtk.simulator.bionet.io_tools import io
    from bmtk.simulator.bionet.iclamp import IClamp, FileIClamp
    from bmtk.simulator.bionet.seclamp import SEClamp
    from bmtk.simulator.bionet import modules as mods
    from bmtk.simulator.core.node_sets import NodeSet
    import bmtk.simulator.utils.simulation_reports as reports
    import bmtk.simulator.utils.simulation_inputs as inputs
    from bmtk.utils.reports.spike_trains import SpikeTrains
    import h5py

    pc = h.ParallelContext()    # object to access MPI methods

    class BioSimulator(Simulator):
        """Includes methods to run and control the simulation"""

        def __init__(self, network, dt, tstop, v_init, celsius, nsteps_block, start_from_state=False):
            self.net = network

            self._start_from_state = start_from_state
            self.dt = dt
            self.tstop = tstop
            self.tstart = 0.0

            self._v_init = v_init
            self._celsius = celsius
            self._h = h

            self.tstep = int(round(h.t / h.dt))
            self.tstep_start_block = self.tstep
            self.nsteps = int(round(h.tstop/h.dt))

            # make sure the block size isn't small than the total number of steps
            # TODO: should we send a warning that block-step size is being reset?
            self._nsteps_block = nsteps_block if self.nsteps > nsteps_block else self.nsteps

            self.__tstep_end_block = 0
            self.__tstep_start_block = 0

            h.runStopAt = h.tstop
            h.steps_per_ms = 1/h.dt
            pc.setup_transfer()#Sets up gap junctions.
            self._set_init_conditions()  # call to save state
            h.cvode.cache_efficient(1)

            h.pysim = self  # use this objref to be able to call postFadvance from proc advance in advance.hoc
            self._iclamps = []
            self._f_iclamps = []
            self._seclamps = []

            self._output_dir = 'output'
            self._log_file = 'output/log.txt'

            self._spikes = {}  # for keeping track of different spike times, key of cell gids

            self._cell_variables = []  # location of saved cell variables
            self._cell_vars_dir = 'output/cellvars'

            self._sim_mods = []  # list of modules.SimulatorMod's

        @property
        def dt(self):
            return h.dt

        @dt.setter
        def dt(self, ms):
            h.dt = ms

        @property
        def tstop(self):
            return h.tstop

        @tstop.setter
        def tstop(self, ms):
            h.tstop = ms

        @property
        def v_init(self):
            return self._v_init

        @v_init.setter
        def v_init(self, voltage):
            self._v_init = voltage

        @property
        def celsius(self):
            return self._celsius

        @celsius.setter
        def celsius(self, c):
            self._celsius = c

        @property
        def n_steps(self):
            return int(round(self.tstop/self.dt))

        @property
        def cell_variables(self):
            return self._cell_variables

        @property
        def cell_var_output(self):
            return self._cell_vars_dir

        @property
        def spikes_table(self):
            return self._spikes

        @property
        def nsteps_block(self):
            return self._nsteps_block

        @property
        def h(self):
            return self._h

        @property
        def biophysical_gids(self):
            return self.net.cell_type_maps('biophysical').keys()

        @property
        def local_gids(self):
            # return self.net.get
            return self.net.local_gids

        def simulation_time(self, units='ms'):
            units_lc = units.lower()
            time_ms = self.tstop - self.tstart
            if units_lc == 'ms':
                return time_ms
            elif units_lc == 's':
                return time_ms/1000.0

            else:
                raise AttributeError('Uknown unit type {}'.format(units))

        def __elapsed_time(self, time_s):
            if time_s < 120:
                return '{:.4} seconds'.format(time_s)
            elif time_s < 7200:
                mins, secs = divmod(time_s, 60)
                return '{} minutes, {:.4} seconds'.format(mins, secs)
            else:
                mins, secs = divmod(time_s, 60)
                hours, mins = divmod(mins, 60)
                return '{} hours, {} minutes and {:.4} seconds'.format(hours, mins, secs)

        def _set_init_conditions(self):
            """Set up the initial conditions: either read from the h.SaveState or from config["condidtions"]"""
            pc.set_maxstep(10)
            h.stdinit()
            self.tstep = int(round(h.t/h.dt))
            self.tstep_start_block = self.tstep

            if self._start_from_state:
                # io.read_state()
                io.log_info('Read the initial state saved at t_sim: {} ms'.format(h.t))
            else:
                h.v_init = self.v_init

            h.celsius = self.celsius

        def set_spikes_recording(self):
            for gid, _ in self.net.get_local_cells().items():
                tvec = self.h.Vector()
                gidvec = self.h.Vector()
                pc.spike_record(gid, tvec, gidvec)
                self._spikes[gid] = tvec

        def attach_se_voltage_clamp(self, amplitudes, durations, gids, rs):
            if gids is None or gids=='all':
                gids = self.biophysical_gids

            if isinstance(gids, int):
                gids = [gids]
            elif isinstance(gids, string_types):
                gids = [int(gids)]
            elif isinstance(gids, NodeSet):
                gids = gids.gids()

            all_gids = list(gids).copy()
            gids = list(set(self.local_gids) & set(gids))

            if len(all_gids)!=len(amplitudes):
                if len(amplitudes) != 1:
                    raise AttributeError("SEClamp must either have amplitudes for each gid or only one set of amplitudes. "
                + "gids: "+ str(all_gids) + ',amps: ' + str(amplitudes) + 'durs,' + str(durations))
                elif len(durations) != 1:
                    raise AttributeError("SEClamp must either have durations for each gid or only one set of durations. "
                + "gids: "+ str(all_gids) + ',amps: ' + str(amplitudes) + 'durs,' + str(durations))
                else:
                    amplitudes = list(amplitudes) * len(gids)
                    durations = list(durations) * len(gids)

            for idx,gid in enumerate(all_gids):
                if gid in gids:
                    cell = self.net.get_cell_gid(gid)

                    try:
                        length = len(amplitudes[idx])
                        if length!=3:
                            raise AttributeError("SEClamp amplitudes must be a non-list or a list of length 3.")
                    except:
                        amplitudes[idx] = [amplitudes[idx] for i in range(3)]

                    try:
                        length = len(durations[idx])
                        if length!=3:
                            raise AttributeError("SEClamp durations must be a non-list or a list of length 3.")
                    except:
                        durations[idx] = [durations[idx] for i in range(3)]

                    SEvc = SEClamp(amplitudes[idx], durations[idx], rs=rs[idx])

                    SEvc.attach_current(cell)
                    self._seclamps.append(SEvc)

        def add_mod(self, module):
            self._sim_mods.append(module)

        def run(self):
            """Run the simulation:
            if beginning from a blank state, then will use h.run(),
            if continuing from the saved state, then will use h.continuerun()
            """
            for mod in self._sim_mods:
                if isinstance(mod, mods.ClampReport):
                    if mod.variable == "se":
                        mod.initialize(self, self._seclamps)
                    elif mod.variable == "ic":
                        mod.initialize(self, self._iclamps)
                    elif mod.variable == "f_ic":
                        mod.initialize(self, self._f_iclamps)
                else:
                    mod.initialize(self)

            self.start_time = h.startsw()
            s_time = time.time()
            pc.timeout(0)

            pc.barrier()  # wait for all hosts to get to this point
            io.log_info('Running simulation for {:.3f} ms with the time step {:.3f} ms'.format(self.tstop, self.dt))
            io.log_info('Starting timestep: {} at t_sim: {:.3f} ms'.format(self.tstep, h.t))
            io.log_info('Block save every {} steps'.format(self.nsteps_block))

            if self._start_from_state:
                h.continuerun(h.tstop)
            else:
                h.run(h.tstop)        # <- runs simuation: works in parallel

            pc.barrier()

            for mod in self._sim_mods:
                mod.finalize(self)
            pc.barrier()

            end_time = time.time()

            sim_time = self.__elapsed_time(end_time - s_time)
            io.log_info('Simulation completed in {} '.format(sim_time))

        def report_load_balance(self):
            comptime = pc.step_time()
            avgcomp = pc.allreduce(comptime, 1)/pc.nhost()
            maxcomp = pc.allreduce(comptime, 2)
            io.log_info('Maximum compute time is {} seconds.'.format(maxcomp))
            io.log_info('Approximate exchange time is {} seconds.'.format(comptime - maxcomp))
            if maxcomp != 0.0:
                io.log_info('Load balance is {}.'.format(avgcomp/maxcomp))

        def post_fadvance(self):
            """
            Runs after every execution of fadvance (see advance.hoc)
            Called after every time step to perform computation and save data to memory block or to disk.
            The initial condition tstep=0 is not being saved
            """
            for mod in self._sim_mods:
                mod.step(self, self.tstep)

            self.tstep += 1

            if (self.tstep % self.nsteps_block == 0) or self.tstep == self.nsteps:
                io.log_info('    step:{} t_sim:{:.2f} ms'.format(self.tstep, h.t))
                self.__tstep_end_block = self.tstep
                time_step_interval = (self.__tstep_start_block, self.__tstep_end_block)

                for mod in self._sim_mods:
                    mod.block(self, time_step_interval)

                self.__tstep_start_block = self.tstep   # starting point for the next block

        @classmethod
        def from_config(cls, config, network, set_recordings=True):
            simulation_inputs = inputs.from_config(config)

            # Special case for setting synapses to spontaneously (for a given set of pre-synaptic cell-types). Using this
            # input will change the way the network builds cells/connections and thus needs to be set first.
            for sim_input in simulation_inputs:
                if sim_input.input_type == 'syn_activity':
                    network.set_spont_syn_activity(
                        precell_filter=sim_input.params['precell_filter'],
                        postcell_filter=sim_input.params.get('postcell_filter', {}),
                        timestamps=sim_input.params['timestamps']
                    )

            # The network must be built before initializing the simulator because
            # gap junctions must be set up before the simulation is initialized.
            network.io.log_info('Building cells.')
            network.build_nodes()

            network.io.log_info('Building recurrent connections')
            network.build_recurrent_edges()

            sim = cls(network=network,
                      dt=config.dt,
                      tstop=config.tstop,
                      v_init=config.v_init,
                      celsius=config.celsius,
                      nsteps_block=config.block_step)

            # TODO: Need to create a gid selector
            for sim_input in inputs.from_config(config):
                if sim_input.input_type == 'spikes' and sim_input.module in ['nwb', 'csv', 'sonata', 'h5']:
                    io.log_info('Building virtual cell stimulations for {}'.format(sim_input.name))
                    path = sim_input.params['input_file']
                    spikes = SpikeTrains.load(path=path, file_type=sim_input.module, **sim_input.params)
                    # node_set_opts = sim_input.params.get('node_set', 'all')
                    node_set = network.get_node_set(sim_input.node_set)
                    network.add_spike_trains(
                        spike_trains=spikes,
                        node_set=node_set,
                        spikes_generator=None,
                        sim=sim
                    )

                elif sim_input.input_type == 'spikes' and sim_input.module == 'function':
                    io.log_info('Building virtual cell stimulations for {}'.format(sim_input.name))
                    # path = sim_input.params.get['input_file']
                    spikes_generator = sim_input.params['spikes_function']
                    node_set = network.get_node_set(sim_input.node_set)
                    network.add_spike_trains(
                        spike_trains=None,
                        node_set=node_set,
                        spikes_generator=spikes_generator,
                        sim=sim
                    )

                elif sim_input.module == 'IClamp':
                    sim.add_mod(mods.IClampMod(input_type=sim_input.input_type, **sim_input.params))

                elif sim_input.module == "SEClamp":
                    sim.add_mod(mods.SEClamp(input_type=sim_input.input_type, **sim_input.params))

                elif sim_input.module == 'xstim':
                    sim.add_mod(mods.XStimMod(**sim_input.params))

                elif sim_input.module == 'comsol':
                    sim.add_mod(mods.ComsolMod(**sim_input.params))

                elif sim_input.module == 'syn_activity':
                    pass

                elif sim_input.module == 'replay':
                    io.log_info('Building replay connections "{}"'.format(sim_input.name))
                    spikes = SpikeTrains.load(path=sim_input.params['spikes_file'])
                    network.build_replay_inputs(
                        spike_trains=spikes,
                        edges_path=sim_input.params['edges']['edges_file'],
                        edge_types_path=sim_input.params['edges']['edge_types_file'],
                        source_node_set=sim_input.params.get('source_node_set', 'all'),
                        target_node_set=sim_input.params.get('target_node_set', 'all')
                    )

                elif sim_input.module == 'ecephys_probe':
                    sim.add_mod(mods.BioECEphysUnitsModule(name=sim_input.name, **sim_input.params))

                else:
                    io.log_exception('Can not parse input format {}'.format(sim_input.name))

            # Parse the "reports" section of the config and load an associated output module for each report
            sim_reports = reports.from_config(config)
            for report in sim_reports:
                if isinstance(report, reports.SpikesReport):
                    mod = mods.SpikesMod(**report.params)

                elif report.module == 'netcon_report':
                    mod = mods.NetconReport(**report.params)

                elif isinstance(report, reports.MembraneReport):
                    if report.params['sections'] == 'soma':
                        mod = mods.SomaReport(**report.params)

                    else:
                        mod = mods.MembraneReport(**report.params)
                elif isinstance(report, reports.ClampReport):
                    mod = mods.ClampReport(**report.params)

                elif isinstance(report, reports.ECPReport):
                    mod = mods.EcpMod(**report.params)
                    # Set up the ability for ecp on all relevant cells
                    # TODO: According to spec we need to allow a different subset other than only biophysical cells
                    for gid, cell in network.cell_type_maps('biophysical').items():
                        cell.setup_ecp()

                elif report.module == 'save_synapses':
                    mod = mods.SaveSynapses(**report.params)

                else:
                    # TODO: Allow users to register customized modules using pymodules
                    io.log_warning('Unrecognized module {}, skipping.'.format(report.module))
                    continue

                sim.add_mod(mod)

            return sim