# 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 pandas as pd
import h5py
import numpy as np

from .utils import range_itr, get_attribute_h5
from .node import Node, NodeSet
from .edge import Edge, EdgeSet
from .group import NodeGroup, EdgeGroup


class Population(object):
    def __init__(self, pop_name, pop_group, types_table):
        self._pop_name = pop_name
        self._pop_group = pop_group
        self._types_table = types_table
        self._nrows = 0

        # For storing individual groups
        self._group_map = {}  # grp-id --> h5py.Group object
        self._find_groups()
        self._group_cache = {}  # grp-id --> soneta.io.Group() object

        # Refrences to most of the population's primary dataset
        self._type_id_ds = pop_group[self.type_ids_column]
        self._group_id_ds = pop_group[self.group_id_column]
        self._group_index_ds = pop_group[self.group_index_column]

        self._group_indicies = {}  # grp-id --> list of rows indicies
        self._group_indicies_cache_built = False

    @property
    def name(self):
        """name of current population"""
        return self._pop_name

    @property
    def group_ids(self):
        """List of all group_ids belonging to population"""
        return list(self._group_map.keys())

    @property
    def groups(self):
        """Returns a list of sonata.Group objects"""
        return [self.get_group(name) for name in self._group_map.keys()]

    @property
    def types_table(self):
        return self._types_table

    @property
    def type_ids(self):
        return np.array(self._type_id_ds)

    @property
    def group_id_ds(self):
        return self._group_id_ds

    @property
    def group_index_ds(self):
        return self._group_index_ds

    @property
    def group_id_column(self):
        raise NotImplementedError

    @property
    def group_index_column(self):
        raise NotImplementedError

    @property
    def type_ids_column(self):
        raise NotImplementedError

    def to_dataframe(self):
        """Convert Population to dataframe"""
        raise NotImplementedError

    def get_group(self, group_id):
        if group_id in self._group_cache:
            return self._group_cache[group_id]
        else:
            grp_h5 = self._group_map[group_id]
            grp_obj = self._build_group(group_id, grp_h5)
            self._group_cache[group_id] = grp_obj
            return grp_obj

    def group_indicies(self, group_id, build_cache=False):
        """Returns a list of all the population row index that maps onto the given group.

        Used for iterating or searching within a Group

        :param group_id: id of a given group
        :param build_cache: Will cache indicies for all groups. Will be faster if making multiple calls but requires
         more memory (default False)
        :return: A (possibly empty) list of row indicies (non-contiguous, but unique)
        """
        if self._group_indicies_cache_built:
            return self._group_indicies.get(group_id, [])

        else:
            tmp_index = pd.DataFrame()
            # TODO: Need to check the memory overhead, especially for edges. See if an iterative search is just as fast
            tmp_index['grp_id'] = pd.Series(self._group_id_ds[()], dtype=self._group_id_ds.dtype)
            tmp_index['row_indx'] = pd.Series(range_itr(self._nrows), dtype=np.uint32)
            if build_cache:
                # save all indicies as arrays
                self._group_indicies = {grp_id: np.array(subset['row_indx'])
                                        for grp_id, subset in tmp_index.groupby(by='grp_id')}
                self._group_indicies_cache_built = True
                return self._group_indicies.get(group_id, [])
            else:
                # TODO: Manually del tmp_index to clear out the memory?
                tmp_index = tmp_index[tmp_index['grp_id'] == group_id]
                return np.array(tmp_index['row_indx'])

    def igroup_ids(self, row_indicies):
        return self._group_id_ds[list(row_indicies)]

    def igroup_indicies(self, row_indicies):
        return self._group_index_ds[list(row_indicies)]

    def _find_groups(self):
        """Create a map between group-id and h5py.Group reference"""
        for grp_key, grp_h5 in self._pop_group.items():
            if grp_key.isdigit():
                grp_id = int(grp_key)
                self._group_map[grp_id] = grp_h5
            else:
                # TODO: Should we put a warning if an unrecognized group exists?
                pass

    def _build_group(self, group_id, group_h5):
        raise NotImplementedError

    def __len__(self):
        return self._nrows


class NodePopulation(Population):
    def __init__(self, pop_name, pop_group, node_types_tables):
        super(NodePopulation, self).__init__(pop_name=pop_name, pop_group=pop_group, types_table=node_types_tables)

        # TODO: node_ids can be implicit
        self._node_id_ds = pop_group['node_id']
        self._nrows = len(self._node_id_ds)

        # TODO: This isn't necessary if only using iterator. Delay building index until get_node() is called.
        self._index_nid2row = None  # A lookup from node_id --> h5 row number
        self._node_id_index_built = False
        self._build_node_id_index()

        # indicies for gid <--> node_id map
        self._has_gids = False
        self._index_gid2row = None  # gid --> row (for searching by gid)
        self._index_row2gid = None  # row --> gid (for iterator or searching by node-id)
        self._gid_lookup_fnc = lambda _: None  # for looking up gid by row, use fnc pointer rather than conditional

        self.__itr_index = 0  # for iterator

    @property
    def group_id_column(self):
        return 'node_group_id'

    @property
    def group_index_column(self):
        return 'node_group_index'

    @property
    def type_ids_column(self):
        return 'node_type_id'

    @property
    def has_gids(self):
        return self._has_gids

    @property
    def node_ids(self):
        return np.array(self._node_id_ds)

    @property
    def gids(self):
        if self.has_gids:
            return np.array(self._index_gid2row.index)
        else:
            return None

    @property
    def node_types_table(self):
        return self._types_table

    @property
    def index_column_name(self):
        return 'node_id'

    @property
    def node_types_table(self):
        return self.types_table

    def add_gids(self, gid_map_df, force=False):
        if self.has_gids and not force:
            # TODO: not sure if it's best to return an exception or just continue on in silence?
            raise Exception('Node population {} already has gids mapped onto node-ids.'.format(self.name))
            # return

        # Create map from gid --> node_id --> row #
        self._build_node_id_index()
        tmp_df = pd.DataFrame()
        tmp_df['row_id'] = self._index_nid2row.index
        tmp_df['node_id'] = self._index_nid2row
        gid_map_df = gid_map_df.merge(tmp_df, how='left', left_on='node_id', right_on='node_id')
        gid_map_df = gid_map_df.drop(['node_id', 'population'], axis=1)
        self._index_gid2row = gid_map_df.set_index('gid')
        self._index_row2gid = gid_map_df.set_index('row_id')
        self._gid_lookup_fnc = lambda row_indx: self._index_row2gid.loc[row_indx]['gid']
        self._has_gids = True

    def to_dataframe(self, index_by_id=True):
        if len(self.groups) == 1:
            ret_df = self.get_group(self.group_ids[0]).to_dataframe()
        else:
            ret_df = pd.DataFrame()
            for grp_id in self.group_ids:
                ret_df = ret_df.append(self.get_group(grp_id).to_dataframe(), sort=False)

        if index_by_id:
            ret_df = ret_df.set_index('node_id')

        return ret_df

    def get_row(self, row_indx):
        # TODO: Use helper function so we don't have to lookup gid/node_id twice
        # Note: I'm not cacheing the nodes for memory purposes, but it might be beneificial too.
        node_id = self._node_id_ds[row_indx]
        node_type_id = self._type_id_ds[row_indx]
        node_group_id = self._group_id_ds[row_indx]
        node_group_index = self._group_index_ds[row_indx]

        node_type_props = self.node_types_table[node_type_id]
        node_group_props = self.get_group(node_group_id)[node_group_index]
        node_gid = self._gid_lookup_fnc(row_indx)

        return Node(node_id, self._pop_name, node_type_id, node_type_props, node_group_id, node_group_props, None, gid=node_gid)

    def get_rows(self, row_indicies):
        """Returns a set of all nodes based on list of row indicies.

        Warning: currently due to the use of h5py, the list must be ordered and cannot contain duplicates.

        :param row_indicies: A list of row indicies
        :return: An iterable NodeSet of nodes in the specified indicies
        """
        # TODO: Check that row_indicies is unsigned and the max (which will be the last value) < n_rows
        # TODO: Check order and check for duplicates in list
        return NodeSet(row_indicies, self)

    def inode_ids(self, row_indicies):
        # You get errors if row_indicies is a numpy array or panda series so convert to python list
        # TODO: list conversion can be expensive, see if h5py will work with np arrays natively.
        return self._node_id_ds[list(row_indicies)]

    def igids(self, row_indicies):
        gids = self._gid_lookup_fnc(row_indicies)
        if gids is not None:
            gids = np.array(gids)
        return gids

    def inode_type_ids(self, row_indicies):
        # self._node_type_id_ds
        return self._type_id_ds[list(row_indicies)]

    def get_node_id(self, node_id):
        row_indx = self._index_nid2row.loc[node_id]
        return self.get_row(row_indx)

    def get_gid(self, gid):
        # assert(self.has_gids)
        row_indx = self._index_gid2row.loc[gid]['row_id']
        return self.get_row(row_indx)

    def filter(self, **filter_props):
        for grp in self.groups:
            for node in grp.filter(**filter_props):
                yield node

    def _build_node_id_index(self, force=False):
        if self._node_id_index_built and not force:
            return

        self._index_nid2row = pd.Series(range_itr(self._nrows), index=self._node_id_ds, dtype=self._node_id_ds.dtype)
        self._node_id_index_built = True

    def _build_group(self, group_id, group_h5):
        return NodeGroup(group_id, group_h5, self)

    def __iter__(self):
        self.__itr_index = 0
        return self

    def next(self):
        return self.__next__()

    def __next__(self):
        if self.__itr_index >= self._nrows:
            raise StopIteration

        nxt_node = self.get_row(self.__itr_index)
        self.__itr_index += 1
        return nxt_node

    def __getitem__(self, item):
        if isinstance(item, slice):
            # TODO: Check
            start = item.start if item.start is not None else 0
            stop = item.stop if item.stop is not None else self._nrows
            row_indicies = range_itr(start, stop, item.step)
            return NodeSet(row_indicies, self)

        elif isinstance(item, int):
            return self.get_row(item)

        elif isinstance(item, list):
            return NodeSet(item)
        else:
            print('Unable to get item using {}.'.format(type(item)))


class EdgePopulation(Population):
    class __IndexStruct(object):
        """Class sto store indicies subgroup"""
        # TODO: Use collections.namedtuple
        def __init__(self, lookup_table, edge_table):
            self.lookup_table = lookup_table
            self.edge_table = edge_table

    def __init__(self, pop_name, pop_group, edge_types_tables):
        super(EdgePopulation, self).__init__(pop_name=pop_name, pop_group=pop_group, types_table=edge_types_tables)

        # keep reference to source and target datasets
        self._source_node_id_ds = pop_group['source_node_id']
        self._target_node_id_ds = pop_group['target_node_id']

        self._nrows = len(self._source_node_id_ds)

        # TODO: Throw an error/warning if missing
        self._source_population = EdgePopulation.get_source_population(pop_group)
        self._target_population = EdgePopulation.get_target_population(pop_group)

        self.__itr_index = 0

        # TODO: use a function pointer for get_index so it doesn't have to run a conditional every time
        # TODO: add property and/or property so user can determine what indicies exists.
        self._targets_index = None
        self._has_target_index = False
        self._sources_index = None
        self._has_source_index = False
        self.build_indicies()

    @property
    def group_id_column(self):
        return 'edge_group_id'

    @property
    def group_index_column(self):
        return 'edge_group_index'

    @property
    def type_ids_column(self):
        return 'edge_type_id'

    @property
    def source_population(self):
        return self._source_population

    @property
    def target_population(self):
        return self._target_population

    @staticmethod
    def get_source_population(pop_group_h5):
        return get_attribute_h5(pop_group_h5['source_node_id'], 'node_population', None)

    @staticmethod
    def get_target_population(pop_group_h5):
        return get_attribute_h5(pop_group_h5['target_node_id'], 'node_population', None)

    @property
    def edge_types_table(self):
        return self._types_table

    def to_dataframe(self):
        raise NotImplementedError()


    def build_indicies(self):
        indicies_grp = None
        for grp_name in ['indices', 'indicies']:
            if grp_name in self._pop_group:
                indicies_grp = self._pop_group[grp_name]
                break

        # if 'indicies' in self._pop_group:
        if indicies_grp is not None:
            # indicies_grp = self._pop_group['indicies']
            for index_name, index_grp in indicies_grp.items():
                # TODO: Let __IndexStruct build the indicies
                # Make sure subgroup has the correct datasets
                if not isinstance(index_grp, h5py.Group):
                    continue

                if 'node_id_to_range' not in index_grp:
                    # TODO: make this more general, i.e 'id_to_range' thus we can index on gids, edge_types, etc
                    # TODO: Check that there are two columns in dataset
                    raise Exception('index {} in {} edges is missing column {}.'.format(index_name, self.name,
                                                                                        'node_id_to_range'))
                if 'range_to_edge_id' not in index_grp:
                    raise Exception('index {} in {} edges is missing column {}.'.format(index_name, self.name,
                                                                                        'range_to_edge_id'))

                # Cache the index
                targets_lookup = index_grp['node_id_to_range']
                edges_range = index_grp['range_to_edge_id']
                index_obj = self.__IndexStruct(targets_lookup, edges_range)

                # Determine the type of index
                if index_name == 'source_to_target':
                    self._sources_index = index_obj
                    self._has_source_index = True
                elif index_name == 'target_to_source':
                    self._targets_index = index_obj
                    self._has_target_index = True
                else:
                    # TODO: Need to send this to a logger rather than stdout
                    print('Unrecognized index {}. Ignoring.'.format(index_name))

    def _build_group(self, group_id, group_h5):
        return EdgeGroup(group_id, group_h5, self)

    def group_indicies(self, group_id, build_cache=False, as_list=False):
        if as_list:
            return super(EdgePopulation, self).group_indicies(group_id, build_cache)

        # For nodes it's safe to just keep a list of all indicies that map onto a given group. For edges bc there are
        # many more rows (and typically a lot less groups), We want to build an index like for source/target ids
        if len(self._group_map) == 1:
            return len(self), [[0, len(self)]]

        grp_indicies = super(EdgePopulation, self).group_indicies(group_id, build_cache=False)
        if len(grp_indicies) == 0:
            # Return an index with no ranges
            return 0, []

        # cluster into ranges. Naively implement, there is probably a faster way to cluster an ordered array!
        range_beg = grp_indicies[0]
        ranges = []
        for i in range_itr(1, len(grp_indicies)):
            if (grp_indicies[i-1]+1) != grp_indicies[i]:
                ranges.append([range_beg, grp_indicies[i-1]+1])
                range_beg = grp_indicies[i]
        ranges.append([range_beg, grp_indicies[-1]+1])
        return len(grp_indicies), np.array(ranges, dtype=np.uint32)

    '''
    def _get_target_index(self):
        # TODO: Do only once
        if self._targets_index is not None:
            return self._targets_index

        if 'incidies' in self._pop_group:
            if 'target_to_source' in self._pop_group['incidies']:
                targets_lookup = self._pop_group['incidies']['target_to_source']['node_id_to_range']
                edges_range = self._pop_group['incidies']['target_to_source']['range_to_edge_id']
                self._targets_index = self.__IndexStruct(targets_lookup, edges_range)
                return self._targets_index

        # TODO: What to do if index doesn't exist?
        raise NotImplementedError
    '''

    def get_row(self, index):
        src_node = self._source_node_id_ds[index]
        trg_node = self._target_node_id_ds[index]
        edge_type_id = self._type_id_ds[index]
        edge_types_props = self.edge_types_table[edge_type_id]

        edge_group_id = self._group_id_ds[index]
        edge_group_index = self._group_index_ds[index]
        edge_group_props = self.get_group(edge_group_id)[edge_group_index]
        return Edge(trg_node_id=trg_node, src_node_id=src_node, source_pop=self.source_population,
                    target_pop=self.target_population, group_id = edge_group_id,
                    group_props=edge_group_props, edge_types_props=edge_types_props)

    def filter(self, **filter_props):
        selected_edge_types = set(self.edge_types_table.edge_type_ids)
        types_filter = False  # Do we need to filter results by edge_type_id
        if 'edge_type_id' in filter_props:
            # TODO: Make sure the edge_type_id is valid
            selected_edge_types = set([filter_props['edge_type_id']])
            del filter_props['edge_type_id']
            types_filter = True

        if 'group_id' in filter_props:
            grp_id = filter_props['group_id']
            grp_id = [grp_id] if np.isscalar(grp_id) else grp_id
            selected_groups = set(grp_id)
            del filter_props['group_id']
        else:
            selected_groups = set(self._group_map.keys())  # list of grp_id's that will be used

        group_prop_filter = {}  # list of actual query statements
        group_filter = False  # do we need to filter results by group_id

        # Go through filter key==value pairs, create filters for groups and edge_types
        for filter_key, filter_val in filter_props.items():
            # Find out what groups, if any, the column should search in.
            group_query = False  # If it's querying a group property don't look in edge_types
            types_query = False
            for grp_id, grp_h5 in self._group_map.items():
                if filter_key in grp_h5:
                    # TODO: Need to check the dtype's match
                    selected_groups &= set([grp_id])
                    group_prop_filter[filter_key] = filter_val
                    group_query = True
                    group_filter = True

            if (not group_query) and filter_key in self.edge_types_table.columns:
                # Presearch the edge types and get only those edge_type_ids which match key==val
                selected_edge_types &= set(self.edge_types_table.find(filter_key, filter_val))
                types_filter = True
                types_query = True

            if not (group_query or types_query):
                # Property key neither exists in a group or the edge_types_table
                raise Exception('Could not find property {}'.format(filter_key))

        # Iterate through all nodes, only returning those that match the filter
        for indx in range_itr(self._nrows):
            # Filter by edge_type_id
            if types_filter:
                # TODO: Invert the selected_edge_types, it will be faster to fail immeditely than search the entire list
                if self._type_id_ds[indx] not in selected_edge_types:
                    continue

            # Filter by group properties
            if group_filter:
                # TODO: Invert group search
                grp_id = self._group_id_ds[indx]
                if grp_id not in selected_groups:
                    continue

                grp_index = self._group_index_ds[indx]
                search_failed = True
                for prop_key, prop_val in group_prop_filter.items():
                    if prop_val != self._group_map[grp_id][prop_key][grp_index]:
                        break
                else:
                    search_failed = False

                if search_failed:
                    continue

            yield self.get_row(indx)

    def get_target(self, target_node_id):
        # TODO: Raise an exception, or call find() and log a warning that the index is not available
        # TODO: check validity of target_node_id (non-negative integer and smaller than index range)
        assert(self._has_target_index)
        return self._get_index(self._targets_index, target_node_id)

    def get_targets(self, target_node_ids):
        # TODO: verify input is iterable
        assert(self._has_target_index)
        trg_index = self._targets_index
        for trg_id in target_node_ids:
            for edge in self._get_index(trg_index, trg_id):
                yield edge

    def get_source(self, source_node_id):
        assert(self._has_source_index)
        return self._get_index(self._sources_index, source_node_id)

    def get_sources(self, source_node_ids):
        assert(self._has_target_index)
        trg_index = self._sources_index
        for src_id in source_node_ids:
            for edge in self._get_index(trg_index, src_id):
                yield edge

    def _get_index(self, index_struct, lookup_id):
        # TODO: Use a EdgeSet instead
        if lookup_id >= len(index_struct.lookup_table):
            return []
            # raise StopIteration

        edges_table = index_struct.edge_table
        lookup_beg, lookup_end = index_struct.lookup_table[lookup_id]
        for i in range_itr(lookup_beg, lookup_end):
            edge_indx_beg, edge_indx_end = edges_table[i]
            for edge_indx in range_itr(edge_indx_beg, edge_indx_end):
                yield self.get_row(edge_indx)

    def __iter__(self):
        self.__itr_index = 0
        return self

    def __next__(self):
        if self.__itr_index >= self._nrows:
            raise StopIteration

        next_edge = self.get_row(self.__itr_index)
        self.__itr_index += 1
        return next_edge

    def next(self):
        return self.__next__()
