# Source code for pymatgen.analysis.chemenv.connectivity.structure_connectivity

    """
    Structure connectivity class.
    """

    import networkx as nx
    import numpy as np
    import collections
    from pymatgen.analysis.chemenv.coordination_environments.structure_environments import LightStructureEnvironments
    from pymatgen.analysis.chemenv.connectivity.environment_nodes import get_environment_node
    from pymatgen.analysis.chemenv.connectivity.connected_components import ConnectedComponent
    from monty.json import MSONable
    from monty.json import jsanitize
    import logging

    __author__ = "David Waroquiers"
    __copyright__ = "Copyright 2012, The Materials Project"
    __credits__ = "Geoffroy Hautier"
    __version__ = "1.0"
    __maintainer__ = "David Waroquiers"
    __email__ = "david.waroquiers@gmail.com"
    __date__ = "June 25, 2019"

    def get_delta_image(isite1, isite2, data1, data2):
        """
        Helper method to get the delta image between one environment and another
        from the ligand's delta images.
        """
        if data1['start'] == isite1:
            if data2['start'] == isite2:
                return np.array(data1['delta']) - np.array(data2['delta'])
            else:
                return np.array(data1['delta']) + np.array(data2['delta'])
        else:
            if data2['start'] == isite2:
                return -np.array(data1['delta']) - np.array(data2['delta'])
            else:
                return -np.array(data1['delta']) + np.array(data2['delta'])

    class StructureConnectivity(MSONable):
        """
        Main class containing the connectivity of a structure.
        """

        def __init__(self, light_structure_environment, connectivity_graph=None, environment_subgraphs=None):
            """
            Constructor for the StructureConnectivity object.

            Args:
                light_structure_environment: a LightStructureEnvironments object
                                             containing the relevant local environments
                                             for the sites in the structure.
                connectivity_graph: the networkx MultiGraph if it has already been computed,
                                    e.g. stored in a file or dict and StructureConnectivity
                                    is reconstructed from that file or dict.
                environment_subgraphs: the different subgraphs of environments that have
                                       been computed if any (as for connectivity_graph, only
                                       if it is reconstructed from a file or dict).
            """
            self.light_structure_environments = light_structure_environment
            if connectivity_graph is None:
                self._graph = nx.MultiGraph()
            else:
                self._graph = connectivity_graph
            if environment_subgraphs is None:
                self.environment_subgraphs = {}
            else:
                self.environment_subgraphs = environment_subgraphs

        def environment_subgraph(self, environments_symbols=None, only_atoms=None):
            """

            Args:
                environments_symbols ():
                only_atoms ():

            Returns:

            """
            if environments_symbols is not None:
                self.setup_environment_subgraph(environments_symbols=environments_symbols, only_atoms=only_atoms)
            try:
                return self._environment_subgraph
            except AttributeError:
                all_envs = self.light_structure_environments.environments_identified()
                self.setup_environment_subgraph(environments_symbols=all_envs, only_atoms=only_atoms)
            return self._environment_subgraph

        def add_sites(self):
            """
            Add the sites in the structure connectivity graph.
            """
            self._graph.add_nodes_from(list(range(len(self.light_structure_environments.structure))))

        def add_bonds(self, isite, site_neighbors_set):
            """
            Add the bonds for a given site index to the structure connectivity graph.

            Args:
                isite: Index of the site for which the bonds have to be added.
                site_neighbors_set: site_neighbors_set: Neighbors set of the site
            """
            existing_edges = self._graph.edges(nbunch=[isite], data=True)
            for nb_index_and_image in site_neighbors_set.neighb_indices_and_images:
                nb_index_unitcell = nb_index_and_image['index']
                nb_image_cell = nb_index_and_image['image_cell']
                exists = False
                if np.allclose(nb_image_cell, np.zeros(3)):
                    for (isite1, ineighb1, data1) in existing_edges:
                        if np.allclose(data1['delta'], np.zeros(3)) and nb_index_unitcell == ineighb1:
                            exists = True
                            break
                else:
                    if isite == nb_index_unitcell:
                        for (isite1, ineighb1, data1) in existing_edges:
                            if isite1 == ineighb1:
                                if np.allclose(data1['delta'],
                                               nb_image_cell) or np.allclose(data1['delta'],
                                                                             -nb_image_cell):
                                    exists = True
                                    break
                    else:
                        for (isite1, ineighb1, data1) in existing_edges:
                            if nb_index_unitcell == ineighb1:
                                if data1['start'] == isite:
                                    if np.allclose(data1['delta'], nb_image_cell):
                                        exists = True
                                        break
                                elif data1['end'] == isite:
                                    if np.allclose(data1['delta'], -nb_image_cell):
                                        exists = True
                                        break
                                else:
                                    raise ValueError('SHOULD NOT HAPPEN ???')
                if not exists:
                    self._graph.add_edge(isite, nb_index_unitcell, start=isite, end=nb_index_unitcell, delta=nb_image_cell)

        def setup_environment_subgraph(self, environments_symbols, only_atoms=None):
            """
            Set up the graph for predefined environments and optionally atoms.

            Args:
                environments_symbols: Symbols of the environments for the environment subgraph.
                only_atoms: Atoms to be considered.
            """
            logging.info('Setup of environment subgraph for environments {}'.format(', '.join(environments_symbols)))
            if not isinstance(environments_symbols, collections.abc.Iterable):
                environments_symbols = [environments_symbols]
            environments_symbols = sorted(environments_symbols)
            envs_string = '-'.join(environments_symbols)
            if only_atoms is not None:
                envs_string += '#' + '-'.join(sorted(only_atoms))
            # Get it directly if it was already computed
            if envs_string in self.environment_subgraphs:
                self._environment_subgraph = self.environment_subgraphs[envs_string]
                return

            # Initialize graph for a subset of environments
            self._environment_subgraph = nx.MultiGraph()
            # Add the sites with the required environment(s)
            for isite, ce_this_site_all in enumerate(self.light_structure_environments.coordination_environments):
                if ce_this_site_all is None:
                    continue
                if len(ce_this_site_all) == 0:
                    continue
                ce_this_site = ce_this_site_all[0]['ce_symbol']
                if ce_this_site in environments_symbols:
                    if only_atoms is None:
                        env_node = get_environment_node(self.light_structure_environments.structure[isite], isite,
                                                        ce_this_site)
                        self._environment_subgraph.add_node(env_node)
                    else:
                        if self.light_structure_environments.structure.is_ordered:
                            if self.light_structure_environments.structure[isite].specie.symbol in only_atoms:
                                env_node = get_environment_node(self.light_structure_environments.structure[isite], isite,
                                                                ce_this_site)
                                self._environment_subgraph.add_node(env_node)
                        else:
                            #  TODO: add the possibility of a "constraint" on the minimum percentage
                            #        of the atoms on the site
                            this_site_elements = [sp.symbol for sp in
                                                  self.light_structure_environments.structure[isite].species_and_occu]
                            for elem_symbol in this_site_elements:
                                if elem_symbol in only_atoms:
                                    env_node = get_environment_node(self.light_structure_environments.structure[isite],
                                                                    isite, ce_this_site)
                                    self._environment_subgraph.add_node(env_node)
                                    break
            # Find the connections between the environments
            nodes = list(self._environment_subgraph.nodes())
            for inode1, node1 in enumerate(nodes):
                isite1 = node1.isite
                links_node1 = self._graph.edges(isite1, data=True)
                for inode2, node2 in enumerate(nodes[inode1:]):
                    isite2 = node2.isite
                    links_node2 = self._graph.edges(isite2, data=True)
                    # We look for ligands that are common to both site1 and site2
                    connections_site1_site2 = {}
                    for (site1_1, ilig_site1, d1) in links_node1:
                        for (site2_1, ilig_site2, d2) in links_node2:
                            if ilig_site1 == ilig_site2:
                                delta_image = get_delta_image(isite1, isite2, d1, d2)
                                if isite1 == isite2 and np.all(delta_image == 0):
                                    continue
                                tuple_delta_image = tuple(delta_image)
                                if tuple_delta_image in connections_site1_site2:
                                    connections_site1_site2[tuple_delta_image].append((ilig_site1, d1, d2))
                                else:
                                    connections_site1_site2[tuple_delta_image] = [(ilig_site1, d1, d2)]
                    # Remove the double self-loops ...
                    if isite1 == isite2:
                        remove_deltas = []
                        alldeltas = list(connections_site1_site2.keys())
                        alldeltas2 = list(connections_site1_site2.keys())
                        if (0, 0, 0) in alldeltas:
                            alldeltas.remove((0, 0, 0))
                            alldeltas2.remove((0, 0, 0))
                        for current_delta in alldeltas:
                            opp_current_delta = tuple([-dd for dd in current_delta])
                            if opp_current_delta in alldeltas2:
                                remove_deltas.append(current_delta)
                                alldeltas2.remove(current_delta)
                                alldeltas2.remove(opp_current_delta)
                        for remove_delta in remove_deltas:
                            connections_site1_site2.pop(remove_delta)
                    # Add all the edges
                    for conn, ligands in list(connections_site1_site2.items()):
                        self._environment_subgraph.add_edge(node1, node2, start=node1.isite, end=node2.isite,
                                                            delta=conn, ligands=ligands)
            self.environment_subgraphs[envs_string] = self._environment_subgraph

        def setup_connectivity_description(self):
            """

            Returns:

            """
            pass

        def get_connected_components(self, environments_symbols=None, only_atoms=None):
            """
            Args:
                environments_symbols ():
                only_atoms ():

            Returns:

            """
            connected_components = []
            env_subgraph = self.environment_subgraph(environments_symbols=environments_symbols, only_atoms=only_atoms)
            for component_nodes in nx.connected_components(env_subgraph):
                graph = env_subgraph.subgraph(component_nodes).copy()
                connected_components.append(ConnectedComponent.from_graph(graph))
            return connected_components

        def setup_atom_environment_subgraph(self, atom_environment):
            """

            Args:
                atom_environment ():

            Returns:

            """
            raise NotImplementedError()

        def setup_environments_subgraph(self, environments_symbols):
            """

            Args:
                environments_symbols ():

            Returns:

            """
            raise NotImplementedError()

        def setup_atom_environments_subgraph(self, atoms_environments):
            """

            Args:
                atoms_environments ():

            Returns:

            """
            raise NotImplementedError()

        def print_links(self):
            """

            Returns:

            """
            nodes = self.environment_subgraph().nodes()
            print('Links in graph :')
            for node in nodes:
                print(node.isite, ' is connected with : ')
                for (n1, n2, data) in self.environment_subgraph().edges(node, data=True):
                    if n1.isite == data['start']:
                        print('  - {:d} by {:d} ligands ({:d} {:d} {:d})'.format(n2.isite, len(data['ligands']),
                                                                                 data['delta'][0], data['delta'][1],
                                                                                 data['delta'][2]))
                    else:
                        print('  - {:d} by {:d} ligands ({:d} {:d} {:d})'.format(n2.isite, len(data['ligands']),
                                                                                 -data['delta'][0], -data['delta'][1],
                                                                                 -data['delta'][2]))

        def as_dict(self):
            """

            Returns:

            """
            return {"@module": self.__class__.__module__,
                    "@class": self.__class__.__name__,
                    "light_structure_environments": self.light_structure_environments.as_dict(),
                    "connectivity_graph": jsanitize(nx.to_dict_of_dicts(self._graph)),
                    "environment_subgraphs": {env_key: jsanitize(nx.to_dict_of_dicts(subgraph))
                                              for env_key, subgraph in self.environment_subgraphs.items()}}

        @classmethod
        def from_dict(cls, d):
            """

            Args:
                d ():

            Returns:

            """
            # Reconstructs the graph with integer as nodes (json's as_dict replaces integer keys with str keys)
            cgraph = nx.from_dict_of_dicts(d['connectivity_graph'], create_using=nx.MultiGraph, multigraph_input=True)
            cgraph = nx.relabel_nodes(cgraph, int)  # Just relabel the nodes using integer casting (maps str->int)
            # Relabel multiedges (removes multiedges with str keys and adds them back with int keys)
            edges = set(cgraph.edges())
            for n1, n2 in edges:
                new_edges = {int(iedge): edata for iedge, edata in cgraph[n1][n2].items()}
                cgraph.remove_edges_from([(n1, n2, iedge) for iedge, edata in cgraph[n1][n2].items()])
                cgraph.add_edges_from([(n1, n2, iedge, edata) for iedge, edata in new_edges.items()])
            return cls(LightStructureEnvironments.from_dict(d['light_structure_environments']),
                       connectivity_graph=cgraph,
                       environment_subgraphs=None)
            # TODO: also deserialize the environment_subgraphs
            #            environment_subgraphs={env_key: nx.from_dict_of_dicts(subgraph, multigraph_input=True)
            #                                   for env_key, subgraph in d['environment_subgraphs'].items()})