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

from .network import Network
from bmtk.builder.node import Node
from bmtk.builder.edge import Edge
from bmtk.utils import sonata

from .edges_collator import EdgesCollator
from .edge_props_table import EdgeTypesTable
from ..index_builders import create_index_in_memory, create_index_on_disk
from ..builder_utils import mpi_rank, mpi_size, barrier
from ..edges_sorter import sort_edges


logger = logging.getLogger(__name__)


class DenseNetwork(Network):
    def __init__(self, name, **network_props):
        super(DenseNetwork, self).__init__(name, **network_props or {})
        # self.__edges_types = {}
        # self.__src_mapping = {}
        # self.__networks = {}
        # self.__node_count = 0
        self._nodes = []
        self.__edges_tables = []
        self._target_networks = {}

    def _initialize(self):
        self.__id_map = []
        self.__lookup = []

    def _add_nodes(self, nodes):
        self._nodes.extend(nodes)
        self._nnodes = len(self._nodes)

    def edges_table(self):
        return self.__edges_tables

    def _save_nodes(self, nodes_file_name, compression='gzip'):
        if not self._nodes_built:
            self._build_nodes()
        if compression == 'none':
            compression = None  # legit option for h5py for no compression

        # save the node_types file
        group_indx = 0
        groups_lookup = {}
        group_indicies = {}
        group_props = {}
        for ns in self._node_sets:
            if ns.params_hash in groups_lookup:
                continue
            else:
                groups_lookup[ns.params_hash] = group_indx
                group_indicies[group_indx] = 0
                group_props[group_indx] = {k: [] for k in ns.params_keys if k != 'node_id'}
                group_indx += 1

        node_gid_table = np.zeros(self._nnodes)  # todo: set dtypes
        node_type_id_table = np.zeros(self._nnodes)
        node_group_table = np.zeros(self._nnodes)
        node_group_index_tables = np.zeros(self._nnodes)

        for i, node in enumerate(self.nodes()):
            node_gid_table[i] = node.node_id
            node_type_id_table[i] = node.node_type_id
            group_id = groups_lookup[node.params_hash]
            node_group_table[i] = group_id
            node_group_index_tables[i] = group_indicies[group_id]
            group_indicies[group_id] += 1

            group_dict = group_props[group_id]
            for key, prop_ds in group_dict.items():
                prop_ds.append(node.params[key])

        if mpi_rank == 0:
            with h5py.File(nodes_file_name, 'w') as hf:
                # Add magic and version attribute
                add_hdf5_attrs(hf)

                pop_grp = hf.create_group('/nodes/{}'.format(self.name))
                pop_grp.create_dataset('node_id', data=node_gid_table, dtype='uint64', compression=compression)
                pop_grp.create_dataset('node_type_id', data=node_type_id_table, dtype='uint64', compression=compression)
                pop_grp.create_dataset('node_group_id', data=node_group_table, dtype='uint32', compression=compression)
                pop_grp.create_dataset('node_group_index', data=node_group_index_tables, dtype='uint64', compression=compression)

                for grp_id, props in group_props.items():
                    model_grp = pop_grp.create_group('{}'.format(grp_id))

                    for key, dataset in props.items():
                        try:
                            model_grp.create_dataset(key, data=dataset, compression=compression)
                        except TypeError:
                            str_list = [str(d) for d in dataset]
                            hf.create_dataset(key, data=str_list, compression=compression)
        barrier()

    def nodes_iter(self, node_ids=None):
        if node_ids is not None:
            return [n for n in self._nodes if n.node_id in node_ids]
        else:
            return self._nodes

    def _process_nodepool(self, nodepool):
        return nodepool

    def import_nodes(self, nodes_file_name, node_types_file_name, population=None):
        sonata_file = sonata.File(data_files=nodes_file_name, data_type_files=node_types_file_name)
        if sonata_file.nodes is None:
            raise Exception('nodes file {} does not have any nodes.'.format(nodes_file_name))

        populations = sonata_file.nodes.populations
        if len(populations) == 1:
            node_pop = populations[0]
        elif population is None:
            raise Exception('The nodes file {} contains multiple populations.'.format(nodes_file_name) +
                            'Please specify population parameter.')
        else:
            for pop in populations:
                if pop.name == population:
                    node_pop = pop
                    break
            else:
                raise Exception('Nodes file {} does not contain population {}.'.format(nodes_file_name, population))

        for node_type_props in node_pop.node_types_table:
            self._add_node_type(node_type_props)

        for node in node_pop:
            self._node_id_gen.remove_id(node.node_id)
            self._nodes.append(Node(node.node_id, node.group_props, node.node_type_properties))

    def _add_edges(self, connection_map, i):
        """

        :param connection_map:
        :param i:
        """
        edge_type_id = connection_map.edge_type_properties['edge_type_id']
        logger.debug('Generating edges data for edge_types_id {}.'.format(edge_type_id))
        edges_table = EdgeTypesTable(connection_map, network_name=self.name)
        connections = connection_map.connection_itr()

        # iterate through all possible SxT source/target pairs and use the user-defined function/list/value to update
        # the number of syns between each pair. TODO: See if this can be vectorized easily.
        for conn in connections:
            if conn[2]:
                edges_table.set_nsyns(source_id=conn[0], target_id=conn[1], nsyns=conn[2])

        target_net = connection_map.target_nodes
        self._target_networks[target_net.network_name] = target_net.network

        # For when the user specified individual edge properties to be put in the hdf5 (syn_weight, syn_location, etc),
        # get prop value and add it to the edge-types table. Need to fetch and store SxTxN value (where N is the avg
        # num of nsyns between each source/target pair) and it is necessary that the nsyns table be finished.
        for param in connection_map.params:
            rule = param.rule
            rets_multiple_vals = isinstance(param.names, (list, tuple, np.ndarray))

            if not rets_multiple_vals:
                prop_name = param.names  # name of property
                prop_type = param.dtypes.get(prop_name, None)
                edges_table.create_property(prop_name=param.names, prop_type=prop_type)  # initialize property array

                for source_node, target_node, edge_index in edges_table.iter_edges():
                    # calls connection map rule and saves value to edge table
                    pval = rule(source_node, target_node)
                    edges_table.set_property_value(prop_name=prop_name, edge_index=edge_index, prop_value=pval)

            else:
                # Same as loop above, but some connection-map 'rules' will return multiple properties for each edge.
                pnames = param.names
                ptypes = [param.dtypes[pn] for pn in pnames]
                for prop_name, prop_type in zip(pnames, ptypes):
                    edges_table.create_property(prop_name=prop_name, prop_type=prop_type)  # initialize property arrays

                for source_node, target_node, edge_index in edges_table.iter_edges():
                    pvals = rule(source_node, target_node)
                    for pname, pval in zip(pnames, pvals):
                        edges_table.set_property_value(prop_name=pname, edge_index=edge_index, prop_value=pval)

        logger.debug('Edge-types {} data built with {} connection ({} synapses)'.format(
            edge_type_id, edges_table.n_edges, edges_table.n_syns)
        )

        edges_table.save()

        # To EdgeTypesTable the number of synaptic/gap connections between all source/target paris, which can be more
        # than the number of actual edges stored (for efficency), may be a better user-representation.
        self._nedges += edges_table.n_syns  # edges_table.n_edges
        self.__edges_tables.append(edges_table)

    def _get_edge_group_id(self, params_hash):
        return int(params_hash)

    def _save_gap_junctions(self, gj_file_name, compression='gzip'):
        source_ids = []
        target_ids = []
        src_gap_ids = []
        trg_gap_ids = []

        if compression == 'none':
            compression = None  # legit option for h5py for no compression

        for et in self.__edges_tables:
            try:
                is_gap = et.edge_type_properties['is_gap_junction']
            except:
                continue
            if is_gap:
                if et.source_network != et.target_network:
                    raise Exception("All gap junctions must be two cells in the same network builder.")
                table = et.to_dataframe()
                for index, row in table.iterrows():
                    for _ in range(row["nsyns"]):
                        source_ids.append(row["source_node_id"])
                        target_ids.append(row["target_node_id"])
                        src_gap_ids.append(self._gj_id_gen.next())
                        trg_gap_ids.append(self._gj_id_gen.next())
            else:
                continue

        if len(source_ids) > 0:
            with h5py.File(gj_file_name, 'w') as f:
                add_hdf5_attrs(f)
                f.create_dataset('source_ids', data=np.array(source_ids), compression=compression)
                f.create_dataset('target_ids', data=np.array(target_ids), compression=compression)
                f.create_dataset('src_gap_ids', data=np.array(src_gap_ids), compression=compression)
                f.create_dataset('trg_gap_ids', data=np.array(trg_gap_ids), compression=compression)

    def _save_edges(self, edges_file_name, src_network, trg_network, pop_name=None, sort_by='target_node_id',
                    index_by=('target_node_id', 'source_node_id'), compression='gzip'):
        barrier()

        if compression == 'none':
            compression = None  # legit option for h5py for no compression

        if mpi_rank == 0:
            logger.debug('Saving {} --> {} edges to {}.'.format(src_network, trg_network, edges_file_name))

        filtered_edge_types = [
            # Some edges may not match the source/target population
            et for et in self.__edges_tables
            if et.source_network == src_network and et.target_network == trg_network
        ]

        merged_edges = EdgesCollator(filtered_edge_types, network_name=self.name)
        merged_edges.process()
        n_total_conns = merged_edges.n_total_edges
        barrier()

        if n_total_conns == 0:
            if mpi_rank == 0:
                logger.warning('Was not able to generate any edges using the "connection_rule". Not saving.')
            return

        # Try to sort before writing file, If edges are split across ranks/files for MPI/size issues then we need to
        # write to disk first then sort the hdf5 file
        sort_on_disk = False
        edges_file_name_final = edges_file_name
        if sort_by:
            if merged_edges.can_sort:
                merged_edges.sort(sort_by=sort_by)
            else:
                sort_on_disk = True
                edges_file_name_final = edges_file_name

                edges_file_basename = os.path.basename(edges_file_name)
                edges_file_dirname = os.path.dirname(edges_file_name)
                edges_file_name = os.path.join(edges_file_dirname, '.unsorted.{}'.format(edges_file_basename))
                if mpi_rank == 0:
                    logger.debug('Unable to sort edges in memory, will temporarly save to {}'.format(edges_file_name) +
                                 ' before sorting hdf5 file.')
        barrier()

        if mpi_rank == 0:
            logger.debug('Saving {} edges to disk'.format(n_total_conns))
            pop_name = '{}_to_{}'.format(src_network, trg_network) if pop_name is None else pop_name
            with h5py.File(edges_file_name, 'w') as hf:
                # Initialize the hdf5 groups and datasets
                add_hdf5_attrs(hf)
                pop_grp = hf.create_group('/edges/{}'.format(pop_name))

                pop_grp.create_dataset('source_node_id', (n_total_conns,), dtype='uint64', compression=compression)
                pop_grp['source_node_id'].attrs['node_population'] = src_network
                pop_grp.create_dataset('target_node_id', (n_total_conns,), dtype='uint64', compression=compression)
                pop_grp['target_node_id'].attrs['node_population'] = trg_network
                pop_grp.create_dataset('edge_group_id', (n_total_conns,), dtype='uint16', compression=compression)
                pop_grp.create_dataset('edge_group_index', (n_total_conns,), dtype='uint32', compression=compression)
                pop_grp.create_dataset('edge_type_id', (n_total_conns,), dtype='uint32', compression=compression)

                for group_id in merged_edges.group_ids:
                    # different model-groups will have different datasets/properties depending on what edge information
                    # is being saved for each edges
                    model_grp = pop_grp.create_group(str(group_id))
                    for prop_mdata in merged_edges.get_group_metadata(group_id):
                        model_grp.create_dataset(prop_mdata['name'], shape=prop_mdata['dim'], dtype=prop_mdata['type'], compression=compression)

                # Uses the collated edges (eg combined edges across all edge-types) to actually write the data to hdf5,
                # potentially in multiple chunks. For small networks doing it this way isn't very effiecent, however
                # this has the benefits:
                #  * For very large networks it won't always be possible to store all the data in memory.
                #  * When using MPI/multi-node the chunks can represent data from different ranks.
                for chunk_id, idx_beg, idx_end in merged_edges.itr_chunks():
                    pop_grp['source_node_id'][idx_beg:idx_end] = merged_edges.get_source_node_ids(chunk_id)
                    pop_grp['target_node_id'][idx_beg:idx_end] = merged_edges.get_target_node_ids(chunk_id)
                    pop_grp['edge_type_id'][idx_beg:idx_end] = merged_edges.get_edge_type_ids(chunk_id)
                    pop_grp['edge_group_id'][idx_beg:idx_end] = merged_edges.get_edge_group_ids(chunk_id)
                    pop_grp['edge_group_index'][idx_beg:idx_end] = merged_edges.get_edge_group_indices(chunk_id)

                    for group_id, prop_name, grp_idx_beg, grp_idx_end in merged_edges.get_group_data(chunk_id):
                        prop_array = merged_edges.get_group_property(prop_name, group_id, chunk_id)
                        pop_grp[str(group_id)][prop_name][grp_idx_beg:grp_idx_end] = prop_array

            if sort_on_disk:
                logger.debug('Sorting {} by {} to {}'.format(edges_file_name, sort_by, edges_file_name_final))
                sort_edges(
                    input_edges_path=edges_file_name,
                    output_edges_path=edges_file_name_final,
                    edges_population='/edges/{}'.format(pop_name),
                    sort_by=sort_by,
                    compression=compression,
                    # sort_on_disk=True,
                )
                try:
                    logger.debug('Deleting intermediate edges file {}.'.format(edges_file_name))
                    os.remove(edges_file_name)
                except OSError as e:
                    logger.warning('Unable to remove intermediate edges file {}.'.format(edges_file_name))

            if index_by:
                index_by = index_by if isinstance(index_by, (list, tuple)) else [index_by]
                for index_type in index_by:
                    logger.debug('Creating index {}'.format(index_type))
                    create_index_in_memory(
                        edges_file=edges_file_name_final,
                        edges_population='/edges/{}'.format(pop_name),
                        index_type=index_type,
                        compression=compression
                    )

        barrier()
        del merged_edges

        if mpi_rank == 0:
            logger.debug('Saving completed.')

    def _clear(self):
        self._nedges = 0
        self._nnodes = 0

    def edges_iter(self, trg_gids, src_network=None, trg_network=None):
        matching_edge_tables = self.__edges_tables
        if trg_network is not None:
            matching_edge_tables = [et for et in self.__edges_tables if et.target_network == trg_network]

        if src_network is not None:
            matching_edge_tables = [et for et in matching_edge_tables if et.source_network == src_network]

        for edge_type_table in matching_edge_tables:
            et_df = edge_type_table.to_dataframe()
            et_df = et_df[et_df['target_node_id'].isin(trg_gids)]
            if len(et_df) == 0:
                continue

            edge_type_props = edge_type_table.edge_type_properties
            for row in et_df.to_dict(orient='records'):
                yield Edge(
                    src_gid=row['source_node_id'],
                    trg_gid=row['target_node_id'],
                    edge_type_props=edge_type_props,
                    syn_props=row
                )

    @property
    def nnodes(self):
        if not self.nodes_built:
            return 0
        return self._nnodes

    @property
    def nedges(self):
        return self._nedges


def add_hdf5_attrs(hdf5_handle):
    # TODO: move this as a utility function
    hdf5_handle['/'].attrs['magic'] = np.uint32(0x0A7A)
    hdf5_handle['/'].attrs['version'] = [np.uint32(0), np.uint32(1)]
