ALLEN INSTITUTE

# Source code for bmtk.simulator.core.modules.ecephys_module

    from pathlib import Path
    from typing import Any
    import numpy as np
    import pandas as pd

    from .simulator_module import SimulatorMod
    from bmtk.simulator.core.io_tools import io
    from bmtk.utils import lazy_property

    try:
        import pynwb
        has_pynwb = True
    except ImportError as ie:
        has_pynwb = False

    try:
        from mpi4py import MPI
        comm = MPI.COMM_WORLD
        bcast = comm.bcast
        MPI_rank = comm.Get_rank()
        MPI_size = comm.Get_size()
        has_mpi = True
    except:
        MPI_rank = 0
        MPI_size = 1
        bcast = lambda v, n: v
        has_mpi = False

    try:
        # The ndx-aibs extensions for nwb are required to load the neuropixel data modules,
        # warning: can take some time to load
        file_dir = Path(__file__).parent
        namespace_path = (file_dir/"ndx-aibs-ecephys.namespace.yaml").resolve()
        pynwb.load_namespaces(str(namespace_path))
    except:
        io.log_debug('ECEphysUnitsModule: Unable to load ndx-aibs-ecephys.namespace.yaml')

    class ECEphysUnitsModule(SimulatorMod):
        """
        TODO:
        - Have option to specify the nwb file units and/or get them from the NWB
        - Have option to save units-node mapping to output folder
        """
        def __init__(self, name, **kwargs):
            self._name = name
            self._node_set = kwargs['node_set']

            if not has_pynwb:
                io.log_exception('ECEphysUnitsModule: pynwb is not installed (pip install pynwb), unable to use module.')

            # Load a Strategy for mapping SONATA node_ids to NWB unit_ids
            self._mapping_name = kwargs.get('mapping', 'invalid_strategy').lower()
            if self._mapping_name in ['units_map']:
                self._mapping_strategy = UnitIdMapStrategy(**kwargs)

            elif self._mapping_name in ['sample', 'sample_without_replacement']:
                self._mapping_strategy = SamplingStrategy(with_replacement=False, **kwargs)

            elif self._mapping_name in ['sample_with_replacement']:
                self._mapping_strategy = SamplingStrategy(with_replacement=True, **kwargs)

            else:
                io.log_exception('ECEphysUnitsModule: Invalid "mapping" parameters, options: units_map, sample, sample_with_replacement')

        def initialize(self, sim):
            raise NotImplementedError()

    class NWBFileWrapper(object):
        # A Simple wrapper class for nwb files, mainly to keep track of file-path which I can't get from pynwb
        # TODO: Implement a Singleton so that the same nwb file isn't loaded multiple times
        def __init__(self, nwb_path):
            if isinstance(nwb_path, pynwb.file.NWBFile):
                self._id = nwb_path.identifier
                self._io = nwb_path
            else:
                self._id = nwb_path
                self._io = pynwb.NWBHDF5IO(nwb_path, 'r').read()

        @property
        def uuid(self):
            return self._id

        def __getattr__(self, name):
            return getattr(self.__dict__['_io'], name)

    class TimeWindow(object):
        """
        A class for dealing with different strategies for storing and intrepreting time windows intervals [start, stop],
        mainly for use in filtering spike times. Including converting between seconds/miliseconds, parsing an NWB stimulus
        table, and look up for individual unit time, and dealing with defaults. To initialize::

            tw = TimeWindow(defaults=[interval1, interval2, ...], nwb_files=[session1.nwb, session2.nwb, ...])

        Where interval<i> is a time-window associated with all unit spikes in session<i>.nwb, and can include None (in which
        case it will not filter unit spike times).

        If individual units will have unique time intervals then you can pass in a pandas DataFrame with columns
        [unit_ids, start_times, stop_times], with values in ms::

            tw.units_lu = units_times_table_df

        And to fetch the time window associated with unit, for example unit_id 9999 in session_0.nwb, then call::

            window = tw[9999, 'session_0.nwb']

        If will first check to see if unit 9999 has a special [start_time, stop_time] in units_lu table, and if not then fall
        back to the default for 'session_0.nwb', and in seconds. If unit_id/default_session is then it returns a None value.
        """
        def __init__(self, defaults=None, nwb_files=None):
            self._units_lu = None
            self._default_windows = None
            self.conversion_factor = 1/1000.0

            # By requestion, the "time_window" option can
            #  - None
            #  - a single window; "time_window": [100, 200]
            #  - A stim_table filter: {stim_name: gratings, ori: 90.0, tf: 2.0, ...}
            #  - one unique window for each nwb_ids/files; "time_window": [[0, 100], [300, 400], ...]
            # To handle this we will 1) convert each possible option into a list of 0, 1, or more windows. 2) Check
            # that the num of time windows makes sense for the number of nwb_ids/files. And 3) Create a map of
            # defaults for each nwb_ids/file
            if defaults is not None:
                time_windows = self._tolist(defaults)
                n_windows = len(time_windows)
                if n_windows > 1 and len(nwb_files) != n_windows:
                    # There can be no default time window, or one default for all nwb_files, otherwise the number
                    # of time_windows must correspond to the number of nwb_files
                    io.log_exception('ECEphysUnitsModule: Cannot match each "time_window" with the "input_file"s.')

                if n_windows == 1:
                    # convert [interal] -> [interval, interval, interval, ...]
                    time_windows = time_windows*n_windows
                self._default_windows = {nwb.uuid: self._parse_tw(tw, nwb) for tw, nwb in zip(time_windows, nwb_files)}

        @property
        def units_lu(self):
            return self._units_lu

        @units_lu.setter
        def units_lu(self, units_table):
            if 'start_times' in units_table and 'stop_times' in units_table:
                units_table['start_times'] = units_table['start_times']*self.conversion_factor
                units_table['stop_times'] = units_table['stop_times']*self.conversion_factor

                if 'unit_ids' in units_table.columns:
                    units_table = units_table.set_index('unit_ids')

                self._units_lu = units_table

        def _tolist(self, window):
            if isinstance(window, dict):
                # Is a stimulus_table filter, ex. {"interval_name": "gratings", "ori": 90.0, ...}
                return [window]
            elif isinstance(window, (tuple, list)) and len(window) == 0:
                # Is an empyt list
                return []
            elif isinstance(window, (tuple, list)) and isinstance(window[0], (tuple, list, dict, np.ndarray)):
                # is a list of intervals, ex. [[0.0, 100.0], [200.0, 300.0], {stim:gratings} ...]
                return window
            else:
                # assume is a interval, ex. [0.0, 100.0]
                return [window]

        def _parse_tw(self, interval, nwb_file):
            """Converts intervals, including windows and stim-table filters, into appropiate format [start (s), stop (s)]"""
            if isinstance(interval, dict):
                # If it is a dictionary try to find time interval by filtering on nwb.intervals, eg stimulus_table. The filter
                # is a dictionary {'interval_name': 'flashes', 'col1': val1, 'col2': val2, ...}
                filter = interval.copy()
                stim_name = filter.pop('interval_name', None)
                stim_idx = filter.pop('interval_index', 'all')

                # In the NWB there are separate tables for each stimulus, and sometimes they are stored in the
                # nwb as <flashes>_presentations.
                if stim_name is None:
                    io.log_exception('Stimulus table filter missing "interval_name"')
                if stim_name in nwb_file.intervals.keys():
                    interval_df = nwb_file.intervals[stim_name].to_dataframe()
                elif stim_name + '_presentations' in nwb_file.intervals.keys():
                    interval_df = nwb_file.intervals[stim_name + '_presentations'].to_dataframe()
                else:
                    io.log_exception('interval name "{}" not found in {}'.format(stim_name, nwb_file.uuid))

                # In most cases
                interval_df = filter_table(interval_df, filter)
                if len(interval_df) == 0:
                    return [0.0, np.inf]

                if stim_idx == 'all':
                    start_time = interval_df['start_time'].min()
                    stop_time = interval_df['stop_time'].max()
                else:
                    start_time = interval_df.iloc[stim_idx]['start_time']
                    stop_time = interval_df.iloc[stim_idx]['stop_time']

                # In the NWB stim_table and units_tables uses seconds, do not convert time-window
                return [start_time, stop_time]

            else:
                # Is an interval [stop, start] that is entered in manually in units of ms, convert
                # to seconds so it matches nwb spike_times units (s)
                return [interval[0]/1000.0, interval[1]/1000.0]

        def __getitem__(self, unit_info):
            unit_id, nwb_uuid = unit_info[0], unit_info[1]
            if (self.units_lu is not None) and (unit_id in self.units_lu.index):
                unit = self.units_lu.loc[unit_id]
                return [unit['start_times'], unit['stop_times']]
            elif self._default_windows and nwb_uuid in self._default_windows.keys():
                return self._default_windows[nwb_uuid]
            else:
                return None

    class MappingStrategy(object):
        def __init__(self, **kwargs):
            self._nwb_paths = kwargs['input_file']
            self._filters = kwargs.get('units', {})
            self._simulation_onset = kwargs.get('interval_offset', 0.0)/1000.0
            self._missing_ids = kwargs.get('missing_ids', 'fail')
            self._cache_spike_times = kwargs.get('cache', False)
            self._spike_times_cache = {}

            default_window = kwargs.get('interval', None)
            self._time_window = TimeWindow(defaults=default_window, nwb_files=self.nwb_files)

            self._units_table = None
            self._units2nodes_map = None

        @lazy_property
        def nwb_files(self):
            if not isinstance(self._nwb_paths, (list, tuple)):
                self._nwb_paths = [self._nwb_paths]

            nwb_files = []
            for nwb_path in self._nwb_paths:
                nwb_files.append(NWBFileWrapper(nwb_path))

            return nwb_files

        @property
        def units2nodes_map(self):
            return self._units2nodes_map

        @property
        def units_table(self):
            if self._units_table is None:
                # Combine the units and channels table from the nwb file
                merged_table = None
                for nwb_file in self.nwb_files:
                    units_table = self._load_units_table(nwb_file)
                    units_table = self._filter_units_table(units_table)
                    units_table = units_table[['spike_times']]
                    units_table['nwb_uid'] = nwb_file.uuid
                    merged_table = units_table if merged_table is None else pd.concat((merged_table, units_table))

                # if merged_table is None or len(merged_table) == 0:
                #     io.log_exception('ECEphysUnitsModule: Could not parse units table from nwb_file(s).')

                self._units_table = merged_table

            return self._units_table

        def _load_units_table(self, nwb_file):
            units = nwb_file.units.to_dataframe().reset_index()
            units = units.rename(columns={'id': 'unit_id', 'peak_channel_id': 'channel_id'})
            channels = nwb_file.electrodes.to_dataframe()
            channels = channels.reset_index().rename(columns={'id': 'channel_id'})
            return units.merge(channels, how='left', on='channel_id').set_index('unit_id')

        def _filter_units_table(self, units_table):
            units_table = filter_table(units_table, self._filters)
            return units_table

        def build_map(self, node_set):
            raise NotImplementedError()

        def get_spike_trains(self, node_id, source_population):
            # TODO: Is it worth caching spike-trains so we don't have to do a lookup + filtering
            # every time?
            if node_id not in self._units2nodes_map:
                msg = 'ECEphysUnitsModule: Could not find mapping for node_id {}.'.format(node_id)
                if self._missing_ids == 'fail':
                    io.log_exception(msg)
                elif self._missing_ids == 'warn':
                   io.log_warning(msg)
                   return np.array([])

            unit_id = self._units2nodes_map[node_id]
            spike_times = np.array(self.units_table.loc[unit_id]['spike_times'])
            nwb_uid = self.units_table.loc[unit_id]['nwb_uid']
            time_window = self._time_window[unit_id, nwb_uid]
            if time_window is not None:
                spike_times = spike_times[
                    (time_window[0] <= spike_times) & (spike_times <= time_window[1])
                ]
                spike_times = spike_times - time_window[0] + self._simulation_onset

            spike_times = spike_times*1000.0  # Convert from seconds to miliseconds
            return spike_times

    class UnitIdMapStrategy(MappingStrategy):
        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            try:
                self._mapping_path = kwargs['units']
            except KeyError:
                io.log_exception('ECEphysUnitsModule: Could not find "units" csv path for units-to-nodes mapping.')

        def _filter_units_table(self, units_table):
            return units_table

        def build_map(self, node_set):
            try:
                # TODO: Include population name
                mapping_file_df = self._mapping_path if isinstance(self._mapping_path, pd.DataFrame) else pd.read_csv(self._mapping_path, sep=' ')
                self._units2nodes_map = mapping_file_df[['node_ids', 'unit_ids']].set_index('node_ids').to_dict()['unit_ids']
                if 'start_times' in mapping_file_df:
                    self._time_window.units_lu = mapping_file_df[['unit_ids', 'start_times', 'stop_times']].set_index('unit_ids')

            except (FileNotFoundError, UnicodeDecodeError):
                io.log_exception('ECEphysUnitsModule: {} should be a space separated file with columns "node_ids", "unit_ids"'.format(self._mapping_path))

    class SamplingStrategy(MappingStrategy):
        def __init__(self, with_replacement=False, **kwargs):
            super().__init__(**kwargs)
            self._with_replacement = with_replacement

        def build_map(self, node_set):
            node_ids = node_set.node_ids
            unit_ids = self.units_table.index.values

            # There is no way to randomly map a subset of possible unit_ids to all possible
            # node_ids. Ignore, warn user, or fail depending on config file option
            if (not self._with_replacement) and len(node_ids) > len(unit_ids):
                if self._missing_ids == 'fail':
                    # Fail application
                    io.log_exception('ECEphysUnitsModule: Not enough NWB unit_ids to map onto node_set.')

                # Not all node_ids will have spikes, TODO: Make do this at random?
                node_ids = node_ids[:len(unit_ids)]

            # When running with MPI, need to make sure the sampling of the unit_id maps is
            # consistent across all cores. Shuffle on rank 0 and broadcast the new order to all
            # other ranks
            if MPI_rank == 0:
                shuffled_unit_ids = np.random.choice(
                    unit_ids, size=len(node_ids), replace=self._with_replacement
                )
            else:
                shuffled_unit_ids = None
            shuffled_unit_ids = bcast(shuffled_unit_ids, 0)

            # TODO: Include population name
            # Creates an mapping between SONTA node_ids and NWB unit_ids
            self._units2nodes_map = pd.DataFrame({
                'node_ids': node_ids,
                'unit_ids': shuffled_unit_ids
            }).set_index('node_ids').to_dict()['unit_ids']

    def filter_table(table_df, filters_dict):
        if filters_dict:
            # Filter out only those specified units
            mask = True
            for col_name, filter_val in filters_dict.items():
                try:
                    if isinstance(filter_val, str):
                        mask &= table_df[col_name] == filter_val

                    elif isinstance(filter_val, (list, np.ndarray, tuple)):
                        mask &= table_df[col_name].isin(filter_val)

                    elif isinstance(filter_val, dict):
                        col = filter_val.get('column', col_name)
                        op = filter_val['operation']
                        val = filter_val['value']
                        val = '"{}"'.format(val) if isinstance(val, str) else val
                        expr = 'table_df["{}"] {} {}'.format(col, op, val)
                        mask &= eval(expr)

                    else:
                        mask &= table_df[col_name] == filter_val

                except KeyError as ke:
                    col = filter_val.get('column', col_name) if isinstance(filter_val, dict) else col_name
                    io.log_exception('ECEphysUnitsModule: Could not find "{}" column in units/electrodes table.'.format(col))

            table_df = table_df[mask]

        return table_df