ALLEN INSTITUTE

# Source code for bmtk.utils.reports.spike_trains.plotting

    # Copyright 2020. 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 numpy as np
    import six
    import matplotlib.pyplot as plt
    import types
    import copy
    from functools import partial

    from .spike_trains import SpikeTrains
    from .spike_trains_api import SpikeTrainsAPI

    from matplotlib.ticker import MaxNLocator

    def __get_spike_trains(spike_trains):
        """Make sure SpikeTrainsAPI object is always returned"""
        if isinstance(spike_trains, six.string_types):
            # Load spikes from file
            return SpikeTrains.load(spike_trains)

        elif isinstance(spike_trains, (SpikeTrains, SpikeTrainsAPI)):
            return spike_trains

        raise AttributeError('Could not parse spiketrains. Pass in file-path, SpikeTrains object, or list of the previous')

    def __get_population(spike_trains, population):
        """Helper function to figure out which population of nodes to use."""
        pops = spike_trains.populations
        if population is None:
            # If only one population exists in spikes object/file select that one
            if len(pops) > 1:
                raise Exception('SpikeTrains contains more than one population of nodes. Use "population" parameter '
                                'to specify population to display.')

            else:
                return pops[0]

        elif population not in pops:
            raise Exception('Could not find node population "{}" in SpikeTrains, only found {}'.format(population, pops))

        else:
            return population

    def __get_node_groups(spike_trains, node_groups, population):
        """Helper function for parsing the 'node_groups' params"""
        if node_groups is None:
            # If none are specified by user make a 'node_group' consisting of all nodes
            selected_nodes = spike_trains.node_ids(population=population)
            return [{'node_ids': selected_nodes, 'c': 'b'}], selected_nodes
        else:
            # Fetch all node_ids which can be used to filter the data.
            node_groups = copy.deepcopy(node_groups)  # Make a copy since later we may be altering the dictionary
            selected_nodes = np.array(node_groups[0]['node_ids'])
            for grp in node_groups[1:]:
                if 'node_ids' not in grp:
                    raise AttributeError('Could not find "node_ids" key in node_groups parameter.')
                selected_nodes = np.concatenate((selected_nodes, np.array(grp['node_ids'])))

            return node_groups, selected_nodes

    def plot_raster(spike_trains, with_histogram=True, population=None, node_groups=None, times=None, title=None,
                    show=True, save_as=None, plt_style=None):
        """will create a raster plot (plus optional histogram) from a SpikeTrains object or SONATA Spike-Trains file. Will
        return the figure

        By default will display all nodes, if you want to only display a subset of nodes and/or group together different
        nodes (by node_id) by dot colors and labels then you can use the node_groups, which should be a list of dicts::

            plot_raster('/path/to/my/spike.h5',
            node_groups=[{'node_ids': range(0, 70), 'c': 'b', 'label': 'pyr'},      # first 70 nodes are blue pyr cells
                         {'node_ids': range(70, 100), 'c': 'r', 'label': 'inh'}])   # last 30 nodes are red inh cells

        The histogram will not be grouped.

        :param spike_trains: SpikeTrains object or path to a (SONATA) spikes file.
        :param with_histogram: If True the a histogram will be shown as a small subplot below the scatter plot. Default
            True.
        :param population: string. If a spikes-file contains more than one population of nodes, use this to determine which
            nodes to actually plot. If only one population exists and population=None then the function will find it by
            default.
        :param node_groups: None or list of dicts. Used to group sets of nodes by labels and color. Each grouping should
            be a dictionary with a 'node_ids' key with a list of the ids. You can also add 'label' and 'c' keys for
            label and color. If None all nodes will be labeled and colored the same.
        :param times: (float, float). Used to set start and stop time. If not specified will try to find values from spiking
            data.
        :param title: str, Use to add a title. Default no tile
        :param show: bool to display or not display plot. default True.
        :param save_as: None or str: file-name/path to save the plot as a png/jpeg/etc. If None or empty string will not
            save plot.
        :return: matplotlib figure.Figure object
        """
        if plt_style is not None:
            plt.style.use(plt_style)

        spike_trains = __get_spike_trains(spike_trains=spike_trains)
        pop = __get_population(spike_trains=spike_trains, population=population)
        node_groups, selected_ids = __get_node_groups(spike_trains=spike_trains, node_groups=node_groups, population=pop)

        # Only show a legend if one of the node_groups have an explicit label, otherwise matplotlib will show an empty
        # legend box which looks bad
        show_legend = False

        # Situation where if the last (or first) M nodes don't spike matplotlib will cut off the y range, but it should
        # show these as empty rows. To do this need to keep track of range of all node_ids
        min_id, max_id = np.inf, -1

        spikes_df = spike_trains.to_dataframe(populations=pop, with_population_col=False)
        spikes_df = spikes_df[spikes_df['node_ids'].isin(selected_ids)]
        if times is not None:
            min_ts, max_ts = times[0], times[1]
            spikes_df = spikes_df[(spikes_df['timestamps'] >= times[0]) & (spikes_df['timestamps'] <= times[1])]
        else:
            min_ts = np.min(spikes_df['timestamps'])
            max_ts = np.max(spikes_df['timestamps'])

        # Used to determine
        if with_histogram:
            fig, axes = plt.subplots(2, 1, gridspec_kw={'height_ratios': [7, 1]}, squeeze=True)
            raster_axes = axes[0]
            bottom_axes = hist_axes = axes[1]
        else:
            fig, axes = plt.subplots(1, 1)
            bottom_axes = raster_axes = axes
            hist_axes = None

        for node_grp in node_groups:
            grp_ids = node_grp.pop('node_ids')
            grp_spikes = spikes_df[spikes_df['node_ids'].isin(grp_ids)]

            # If label exists for at-least one group we want to show
            show_legend = show_legend or 'label' in node_grp

            # Finds min/max node_id for all node groups
            min_id = np.min([np.min(grp_ids), min_id])
            max_id = np.max([np.max(grp_ids), max_id])

            raster_axes.scatter(grp_spikes['timestamps'], grp_spikes['node_ids'], lw=1, s=12, marker='|',**node_grp)
            #raster_axes.scatter(grp_spikes['timestamps'], grp_spikes['node_ids'], lw=0, s=8, **node_grp)
            raster_axes.yaxis.set_major_locator(MaxNLocator(integer=True))

        if show_legend:
            raster_axes.legend(loc='upper right', markerscale=1.5)

        if title:
            raster_axes.set_title(title)

        raster_axes.set_ylabel('node_ids')
        raster_axes.set_ylim(min_id - 0.5, max_id + 1)  # add buffering to range else the rows at the ends look cut-off.
        raster_axes.set_xlim(min_ts, max_ts + 1)
        bottom_axes.set_xlabel('timestamps ({})'.format(spike_trains.units(population=pop)))

        if with_histogram:
            # Add a histogram if necessary
            hist_axes.hist(spikes_df['timestamps'], 100)
            hist_axes.set_xlim(min_ts - 0.5, max_ts + 1)
            hist_axes.axes.get_yaxis().set_visible(False)
            raster_axes.set_xticks([])

        if save_as:
            plt.savefig(save_as)

        if show:
            plt.show()

        return fig

    def moving_average(data, window_size=10):
        h = int(window_size / 2)
        x_max = len(data)
        return [np.mean(data[max(0, x - h):min(x_max, x + h)]) for x in range(0, x_max)]

    def plot_rates(spike_trains, population=None, node_groups=None, times=None, smoothing=False,
                   smoothing_params=None, title=None, show=True, save_as=None, plt_style=None):
        """Calculate and plot the rates of each node in a SpikeTrains object or SONATA Spike-Trains file. If start and stop
        times are not specified from the "times" parameter, will try to parse values from the timestamps data.

        If you want to only display a subset of nodes and/or group together different nodes (by node_id) by dot colors and
        labels then you can use the node_groups, which should be a list of dicts::

            plot_rates('/path/to/my/spike.h5',
                        node_groups=[{'node_ids': range(0, 70), 'c': 'b', 'label': 'pyr'},
                                     {'node_ids': range(70, 100), 'c': 'r', 'label': 'inh'}])

        :param spike_trains: SpikeTrains object or path to a (SONATA) spikes file.
        :param population: string. If a spikes-file contains more than one population of nodes, use this to determine which
            nodes to actually plot. If only one population exists and population=None then the function will find it by
            default.
        :param node_groups: None or list of dicts. Used to group sets of nodes by labels and color. Each grouping should
            be a dictionary with a 'node_ids' key with a list of the ids. You can also add 'label' and 'c' keys for
            label and color. If None all nodes will be labeled and colored the same.
        :param times: (float, float). Used to set start and stop time. If not specified will try to find values from spiking
            data.
        :param smoothing: Bool or function. Used to smooth the data. By default (False) no smoothing will be done. If True
            will using a moving average smoothing function. Or use a function pointer.
        :param smoothing_params: dict, parameters when using a function pointer smoothing value.
        :param title: str, Use to add a title. Default no tile
        :param show:  bool to display or not display plot. default True.
        :param save_as: None or str: file-name/path to save the plot as a png/jpeg/etc. If None or empty string will not
            save plot.
        :return: matplotlib figure.Figure object
        """
        if plt_style is not None:
            plt.style.use(plt_style)

        spike_trains = __get_spike_trains(spike_trains=spike_trains)
        pop = __get_population(spike_trains=spike_trains, population=population)
        node_groups, selected_ids = __get_node_groups(spike_trains=spike_trains, node_groups=node_groups, population=pop)

        # Determine if smoothing will be applied to the data
        smoothing_params = smoothing_params or {}  # pass in empty parameters
        if isinstance(smoothing, types.FunctionType):
            smoothing_fnc = partial(smoothing, **smoothing_params)
        elif smoothing:
            smoothing_fnc = partial(moving_average, **smoothing_params)
        else:
            smoothing_fnc = lambda d: d  # Use a filler function that won't do anything

        # get data
        spikes_df = spike_trains.to_dataframe(populations=pop, with_population_col=False)
        spikes_df = spikes_df[spikes_df['node_ids'].isin(selected_ids)]
        if times is not None:
            recording_interval = times[1] - times[0]
            spikes_df = spikes_df[(spikes_df['timestamps'] >= times[0]) & (spikes_df['timestamps'] <= times[1])]
        else:
            recording_interval = np.max(spikes_df['timestamps']) - np.min(spikes_df['timestamps'])

        # Iterate through each group of nodes and add to the same plot
        fig, axes = plt.subplots()
        show_legend = False  # Only show labels if one of the node group has label value
        for node_grp in node_groups:
            show_legend = show_legend or 'label' in node_grp # If label exists for at-least one group we want to show

            grp_ids = node_grp.pop('node_ids')
            grp_spikes = spikes_df[spikes_df['node_ids'].isin(grp_ids)]
            spike_rates = grp_spikes.groupby('node_ids').size() / (recording_interval / 1000.0)
            axes.plot(np.array(spike_rates.index), smoothing_fnc(spike_rates), '.', **node_grp)

        axes.set_ylabel('Firing Rates (Hz)')
        axes.set_xlabel('node_ids')
        if show_legend:
            axes.legend()  # loc='upper right')

        if title:
            axes.set_title(title)

        if save_as:
            plt.savefig(save_as)

        if show:
            plt.show()

        return fig

    def plot_rates_boxplot(spike_trains, population=None, node_groups=None, times=None, title=None, show=True,
                           save_as=None, plt_style=None):
        """Creates a box plot of the firing rates taken from a SpikeTrains object or SONATA Spike-Trains file. If start
        and stop times are not specified from the "times" parameter, will try to parse values from the timestamps data.

        By default will plot all nodes together. To only display a subset of the nodes and/or create groups of nodes use
        the node_groups options::

            plot_rates_boxplot(
                '/path/to/my/spike.h5',
                node_groups=[{'node_ids': range(0, 70), 'label': 'pyr'},
                             {'node_ids': range(70, 100), 'label': 'inh'}]
            )

        :param spike_trains: SpikeTrains object or path to a (SONATA) spikes file.
        :param population: string. If a spikes-file contains more than one population of nodes, use this to determine which
            nodes to actually plot. If only one population exists and population=None then the function will find it by
            default.
        :param node_groups: None or list of dicts. Used to group sets of nodes by labels and color. Each grouping should
            be a dictionary with a 'node_ids' key with a list of the ids. You can also add 'label' and 'c' keys for
            label and color. If None all nodes will be labeled and colored the same.
        :param title: str, Use to add a title. Default no tile
        :param show:  bool to display or not display plot. default True.
        :param save_as: None or str: file-name/path to save the plot as a png/jpeg/etc. If None or empty string will not
            save plot.
        :return: matplotlib figure.Figure object
        """
        if (plt_style is not None):
            plt.style.use(plt_style)

        spike_trains = __get_spike_trains(spike_trains=spike_trains)
        pop = __get_population(spike_trains=spike_trains, population=population)
        node_groups, selected_ids = __get_node_groups(spike_trains=spike_trains, node_groups=node_groups, population=pop)

        spikes_df = spike_trains.to_dataframe(populations=pop, with_population_col=False)
        spikes_df = spikes_df[spikes_df['node_ids'].isin(selected_ids)]
        if times is not None:
            recording_interval = times[1] - times[0]
            spikes_df = spikes_df[(spikes_df['timestamps'] >= times[0]) & (spikes_df['timestamps'] <= times[1])]
        else:
            recording_interval = np.max(spikes_df['timestamps']) - np.min(spikes_df['timestamps'])

        fig, axes = plt.subplots()
        rates_data = []
        rates_labels = []

        if len(node_groups) == 1 and 'label' not in node_groups[0]:
            node_groups[0]['label'] = 'All Nodes'

        for i, node_grp in enumerate(node_groups):
            rates_labels.append(node_grp.get('label', 'Node Group {}'.format(i)))

            grp_ids = node_grp.pop('node_ids')
            grp_spikes = spikes_df[spikes_df['node_ids'].isin(grp_ids)]
            spike_rates = grp_spikes.groupby('node_ids').size() / (recording_interval / 1000.0)

            rates_data.append(spike_rates)

        axes.boxplot(rates_data)
        axes.set_ylabel('Firing Rates (Hz)')
        axes.set_xticklabels(rates_labels)

        if title:
            axes.set_title(title)

        if save_as:
            plt.savefig(save_as)

        if show:
            plt.show()

        return fig