ALLEN INSTITUTE

# Source code for bmtk.analyzer

    # 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 os
    from six import string_types
    import h5py
    import pandas as pd
    import matplotlib.pyplot as plt
    import numpy as np

    from bmtk.utils.sonata.config import SonataConfig as cfg

    def _get_config(config):
        if isinstance(config, string_types):
            return cfg.from_json(config)
        elif isinstance(config, dict):
            return config
        else:
            raise Exception('Could not convert {} (type "{}") to json.'.format(config, type(config)))

    def plot_potential(cell_vars_h5=None, config_file=None, gids=None, show_plot=True, save=False):
        if (cell_vars_h5 or config_file) is None:
            raise Exception('Please specify a cell_vars hdf5 file or a simulation config.')

        if cell_vars_h5 is not None:
            plot_potential_hdf5(cell_vars_h5, gids=gids, show_plot=show_plot,
                                save_as='sim_potential.jpg' if save else None)

        else:
            # load the json file or object
            if isinstance(config_file, string_types):
                config = cfg.from_json(config_file)
            elif isinstance(config_file, dict):
                config = config_file
            else:
                raise Exception('Could not convert {} (type "{}") to json.'.format(config_file, type(config_file)))

            gid_list = gids or config['node_id_selections']['save_cell_vars']
            for gid in gid_list:
                save_as = '{}_v.jpg'.format(gid) if save else None
                title = 'cell gid {}'.format(gid)
                var_h5 = os.path.join(config['output']['cell_vars_dir'], '{}.h5'.format(gid))
                plot_potential_hdf5(var_h5, title, show_plot, save_as)

    def plot_potential_hdf5(cell_vars_h5, gids, title='membrane potential', show_plot=True, save_as=None):
        data_h5 = h5py.File(cell_vars_h5, 'r')
        membrane_trace = data_h5['data']

        time_ds = data_h5['/mapping/time']
        tstart = time_ds[0]
        tstop = time_ds[1]
        x_axis = np.linspace(tstart, tstop, len(membrane_trace), endpoint=True)

        gids_ds = data_h5['/mapping/gids']
        index_ds = data_h5['/mapping/index_pointer']
        index_lookup = {gids_ds[i]: (index_ds[i], index_ds[i+1]) for i in range(len(gids_ds))}
        gids = gids_ds.keys() if gids_ds is None else gids
        for gid in gids:
            var_indx = index_lookup[gid][0]
            plt.plot(x_axis, membrane_trace[:, var_indx], label=gid)

        plt.xlabel('time (ms)')
        plt.ylabel('membrane (mV)')
        plt.title(title)
        plt.legend(markerscale=2, scatterpoints=1)

        if save_as is not None:
            plt.savefig(save_as)

        if show_plot:
            plt.show()

    def plot_calcium(cell_vars_h5=None, config_file=None, gids=None, show_plot=True, save=False):
        if (cell_vars_h5 or config_file) is None:
            raise Exception('Please specify a cell_vars hdf5 file or a simulation config.')

        if cell_vars_h5 is not None:
            plot_calcium_hdf5(cell_vars_h5, gids, show_plot=show_plot, save_as='sim_ca.jpg' if save else None)

        else:
            # load the json file or object
            if isinstance(config_file, string_types):
                config = cfg.from_json(config_file)
            elif isinstance(config_file, dict):
                config = config_file
            else:
                raise Exception('Could not convert {} (type "{}") to json.'.format(config_file, type(config_file)))

            gid_list = gids or config['node_id_selections']['save_cell_vars']
            for gid in gid_list:
                save_as = '{}_v.jpg'.format(gid) if save else None
                title = 'cell gid {}'.format(gid)
                var_h5 = os.path.join(config['output']['cell_vars_dir'], '{}.h5'.format(gid))
                plot_calcium_hdf5(var_h5, title, show_plot, save_as)

    def plot_calcium_hdf5(cell_vars_h5, gids, title='Ca2+ influx', show_plot=True, save_as=None):
        data_h5 = h5py.File(cell_vars_h5, 'r')
        cai_trace = data_h5['cai/data']

        time_ds = data_h5['/mapping/time']
        tstart = time_ds[0]
        tstop = time_ds[1]
        x_axis = np.linspace(tstart, tstop, len(cai_trace), endpoint=True)

        gids_ds = data_h5['/mapping/gids']
        index_ds = data_h5['/mapping/index_pointer']
        index_lookup = {gids_ds[i]: (index_ds[i], index_ds[i+1]) for i in range(len(gids_ds))}
        gids = gids_ds.keys() if gids_ds is None else gids
        for gid in gids:
            var_indx = index_lookup[gid][0]
            plt.plot(x_axis, cai_trace[:, var_indx], label=gid)

        #plt.plot(x_axis, cai_trace)
        plt.xlabel('time (ms)')
        plt.ylabel('calcium [Ca2+]')
        plt.title(title)
        plt.legend(markerscale=2, scatterpoints=1)

        if save_as is not None:
            plt.savefig(save_as)

        if show_plot:
            plt.show()

    def spikes_table(config_file, spikes_file=None):
        config = _get_config(config_file)
        spikes_file = config['output']['spikes_file']
        spikes_h5 = h5py.File(spikes_file, 'r')
        gids = np.array(spikes_h5['/spikes/gids'], dtype=np.uint)
        times = np.array(spikes_h5['/spikes/timestamps'], dtype=float)
        return pd.DataFrame(data={'gid': gids, 'spike time (ms)': times})

        #return pd.read_csv(spikes_ascii, names=['time (ms)', 'cell gid'], sep=' ')

    def nodes_table(nodes_file, population):
        # TODO: Integrate into sonata api
        nodes_h5 = h5py.File(nodes_file, 'r')
        nodes_pop = nodes_h5['/nodes'][population]
        root_df = pd.DataFrame(data={'node_id': nodes_pop['node_id'], 'node_type_id': nodes_pop['node_type_id'],
                                     'node_group_id': nodes_pop['node_group_id'],
                                     'node_group_index': nodes_pop['node_group_index']}) #,
                               #index=[nodes_pop['node_group_id'], nodes_pop['node_group_index']])
        root_df = root_df.set_index(['node_group_id', 'node_group_index'])

        node_grps = np.unique(nodes_pop['node_group_id'])
        for grp_id in node_grps:
            sub_group = nodes_pop[str(grp_id)]
            grp_df = pd.DataFrame()
            for hf_key in sub_group:
                hf_obj = sub_group[hf_key]
                if isinstance(hf_obj, h5py.Dataset):
                    grp_df[hf_key] = hf_obj

            subgrp_len = len(grp_df)
            if subgrp_len > 0:
                grp_df['node_group_id'] = [grp_id]*subgrp_len
                grp_df['node_group_index'] = range(subgrp_len)
                grp_df = grp_df.set_index(['node_group_id', 'node_group_index'])
                root_df = root_df.join(other=grp_df, how='left')

        return root_df.reset_index(drop=True)

    def node_types_table(node_types_file, population):
        return pd.read_csv(node_types_file, sep=' ')