def collapse_nodes(graph, node_map, copy=False):
    """
    Args
    ----
    graph : nx.DiGraph
        Graph with nodes we want to collapse.

    node_map : dict
        A map of existing node names to collapsed names.

    copy : bool(False)
        If True, copy the graph before collapsing the nodes.

    Returns
    -------
    nx.DiGraph
        The graph with the nodes collapsed.
        
    """
    nx.relabel_nodes(graph, node_map, copy=copy)

    # remove any self edges created by the relabeling
    graph.remove_edges_from([(u, v) for u, v in graph.edges_iter()
                             if u == v])

    return graph

def get_dot(G_orig):
    """Change labels and colors to be presented with graphviz"""
    G = G_orig.copy()
    cluster_networks.relabel_graph(G)
    tr = {}
    for node in G.nodes_iter():
        tr[node] = '"{}"'.format(node)
    nx.relabel_nodes(G, tr, copy=False)
    for node in G.nodes_iter():
        label = str(node)
        if len(label) > MAX_LABEL:
            label = u'{}..."'.format(label[:MAX_LABEL])
        G.node[node]['label'] = label
    for node in cluster_networks.get_leaf_nodes(G):
        G.node[node]['color'] = "blue"
    for node in cluster_networks.hybrid_nodes(G):
        G.node[node]['color'] = "#7BFF74"  # light green
        G.node[node]['style'] = 'filled'
    for node in cluster_networks.get_root_nodes(G):
        G.node[node]['color'] = "orange"
    for node in cluster_networks.problematic_treechild_nodes(G):
        G.node[node]['color'] = "#FF77EB"  # light pink
        G.node[node]['style'] = 'filled'
    for u, v in cluster_networks.removable_edges(G):
        G.edge[u][v]['color'] = "red"
    for root in cluster_networks.get_root_nodes(G):
        G.node[root]['label'] = '"R"'
    dot = nx.to_pydot(G).to_string()
    return dot.strip()

def add_edge_from_cnode(G,c,h,l):
    
    #If node is being created at last level
    if l == haplolength - 1:
        
        #Add edge from current node to node with name one higher than existing nodes.
        #Set weight and allele to edge data.
        G.add_edge(c, (level[l+1]+1), weight=frequency[h], allele=haplotype[h][l])
        
    #If node is being created that is not at the last level
    else:
        
        #Define function to relabel node names
        def mapping(x):
            if x > level[l+1]:
                return x+1
            else:
                return x
            
        #Relabel node names which are above the level of the node being added.
        nx.relabel_nodes(G,mapping,copy=False)
            
        #Add edge from current node to node with name one higher than the last node on the next level.
        #Set weight and allele to edge data.
        G.add_edge(c, (level[l+1]+1), weight=frequency[h], allele=haplotype[h][l])

def sub_values(tree, var_values):
    """Substitute given values for variables.

    @param tree: AST

    @type var_values: C{dict}

    @return: AST with L{Var} nodes replaces by
        L{Num}, L{Const}, or L{Bool}
    """
    old2new = dict()
    for u in tree.nodes_iter():
        if u.type != 'var':
            continue
        val = var_values[u.value]
        # instantiate appropriate value type
        if isinstance(val, bool):
            v = nodes.Bool(val)
        elif isinstance(val, int):
            v = nodes.Num(val)
        elif isinstance(val, str):
            v = nodes.Str(val)
        old2new[u] = v
    # replace variable by value
    nx.relabel_nodes(tree, old2new, copy=False)

def createtree(n):
    T=nx.DiGraph()
    slevel=[1]*(len(G.node[n]['haplotype'][0])+1) 
    T.add_node(1)    

    for i, hap in enumerate(G.node[n]['haplotype']):
        cnode = 1
       
        for j, h in enumerate(hap):
            
            if h in [(edata['allele']) for u,v,edata in T.out_edges(cnode,data=True) if 'allele' in edata]:
                T[cnode][v]['weight'] = T[cnode][v]['weight'] + G.node[n]['frequency'][i]
                
                cnode = v
            else:
                     
                if j == len(hap)-1 or slevel[j]==slevel[j+1]:
                    pass
                else:
                    #Define function to relabel node names
                    def mapping(x):
                        if x > slevel[j+1]:
                            return x+1
                        else:
                            return x
            
                    #Relabel node names which are above the level of the node being added.
                    nx.relabel_nodes(T,mapping,copy=False)            

                T.add_edge(cnode, slevel[j+1]+1, weight=G.node[n]['frequency'][i], allele=h)
                cnode = slevel[j+1]+1
                for a in range(j+1,len(slevel)):
                    slevel[a] = slevel[a] + 1
    return T

def coalesce(G, node1, node2):
    """Performs Briggs coalescing. Takes in the graph and two nodes.
    Returns 1 if unable to coalesce, 0 otherwise."""
    # we need to get a modified degrees list for pre-colored nodes
    degs = nx.degree(G)
    for i in range(1, k + 1):
        degs[str(i)] = 9999
    if node1 in G.neighbors(node2) or node2 in G.neighbors(node1):
        print "ERROR:", node1, "and", node2, "share an edge. \
Check your moves list"
        return  # intentionally void to cause error
    elif degs[node1] + degs[node2] >= k:
        print "Failure:", node1, "and", node2, "have combined degree \
= " + str(
            degs[node1] + degs[node2]
        ) + ", which is too high for k =", k
        return 1
    else:
        newedge = []
        for i in range(len(G.neighbors(node2))):
            newedge.append((node1, G.neighbors(node2)[i]))
        G.add_edges_from(newedge)
        G.remove_node(node2)
        nx.relabel_nodes(G, {node1: node1 + node2}, copy=False)
        print "Success:", node1, "and", node2, "have been coalesced"
    return 0

    def __init__(self, *args, basepath=None, **kwargs):
        super(CandidateGraph, self).__init__(*args, **kwargs)
        self.node_counter = 0
        node_labels = {}
        self.node_name_map = {}

        for node_name in self.nodes():
            image_name = os.path.basename(node_name)
            image_path = node_name
            # Replace the default attr dict with a Node object
            self.node[node_name] = Node(image_name, image_path, self.node_counter)

            # fill the dictionary used for relabelling nodes with relative path keys
            node_labels[node_name] = self.node_counter
            # fill the dictionary used for mapping base name to node index
            self.node_name_map[self.node[node_name].image_name] = self.node_counter
            self.node_counter += 1

        nx.relabel_nodes(self, node_labels, copy=False)

        for s, d in self.edges():
            if s > d:
                s, d = d, s
            e = self.edge[s][d]
            e.source = self.node[s]
            e.destination = self.node[d]

        self.creationdate = strftime("%Y-%m-%d %H:%M:%S", gmtime())
        self.modifieddate = strftime("%Y-%m-%d %H:%M:%S", gmtime())

    def __init__(self, *args, basepath=None, **kwargs):
        super(CandidateGraph, self).__init__(*args, **kwargs)
        self.node_counter = 0
        node_labels = {}
        self.node_name_map = {}

        # the node_name is the relative path for the image
        for node_name, node in self.nodes_iter(data=True):
            image_name = os.path.basename(node_name)
            image_path = node_name

            # Replace the default node dict with an object
            self.node[node_name] = Node(image_name, image_path)

            # fill the dictionary used for relabelling nodes with relative path keys
            node_labels[node_name] = self.node_counter
            # fill the dictionary used for mapping base name to node index
            self.node_name_map[self.node[node_name].image_name] = self.node_counter
            self.node_counter += 1

        nx.relabel_nodes(self, node_labels, copy=False)

        # Add the Edge class as a edge data structure
        for s, d, edge in self.edges_iter(data=True):
            self.edge[s][d] = Edge(self.node[s], self.node[d])

def spingraph_from_graph(graph):
    # even_graph = nx.relabel_nodes(graph, lambda x:x*2)
    # odd_graph = nx.relabel_nodes(graph, lambda x:2*x+1)
    # union_graph  = nx.union(even_graph, odd_graph)
    # on the fly union saves about 20% memory, ugly but more efficient
    union_graph  = nx.union(nx.relabel_nodes(graph, lambda x:x*2),nx.relabel_nodes(graph, lambda x:2*x+1))
    # from pudb import set_trace; set_trace()
    for spin_down_node in xrange(1,union_graph.order(),2):
        spin_up_node = spin_down_node -1
        for spin_down_node_neighbour in union_graph[spin_down_node].keys():
            if spin_down_node_neighbour % 2 ==0:
                continue
            if spin_down_node_neighbour < spin_down_node:             # is either top or left neighbour
                if spin_down_node_neighbour == spin_down_node-2:      # is left neighbour
                    union_graph.add_edge(spin_up_node,spin_down_node_neighbour,weight=-p.tso)
                    union_graph.add_edge(spin_down_node_neighbour,spin_up_node,weight=-p.tso)
                else:
                    union_graph.add_edge(spin_up_node,spin_down_node_neighbour,weight=+1j*p.tso)
                    union_graph.add_edge(spin_down_node_neighbour,spin_up_node,weight=-1j*p.tso)
            if spin_down_node_neighbour > spin_down_node:             # is either right or bottom neighbour
                if spin_down_node_neighbour == spin_down_node+2:      # is right neighbour
                    union_graph.add_edge(spin_up_node,spin_down_node_neighbour,weight=p.tso)
                    union_graph.add_edge(spin_down_node_neighbour,spin_up_node,weight=p.tso)
                else:
                    union_graph.add_edge(spin_up_node,spin_down_node_neighbour,weight=-1j*p.tso)
                    union_graph.add_edge(spin_down_node_neighbour,spin_up_node,weight=+1j*p.tso)
    return union_graph

    def get_clusters(self, num_clusters=None):
        if num_clusters == None:
            index, value = max(enumerate(self.quality_history), key=lambda iv: iv[1])
            num_clusters = len(self.quality_history) - index
        num_clusters = max(min(num_clusters, self.max_clusters), 0)

        clusters = [set([n]) for n in self.orphans]
        if self.dendrogram and num_clusters:
            nx.relabel_nodes(self.dendrogram, self.rename_map.integer, copy=False)
            try:
                start_node = max(self.dendrogram)
                priors, fringe = self.dendrogram_crawl(start=start_node, max_fringe_size=num_clusters)

                # Double check we got the right number of values
                if len(fringe) != num_clusters:
                    raise ValueError("Failed to retrieve %d clusters correctly (got %d instead)"
                        % (num_clusters, len(fringe)))

                for neg_clust_start in fringe:
                    clust_start = -neg_clust_start
                    cprior, cfringe = self.dendrogram_crawl(start=clust_start, priors=priors.copy())
                    cluster_set = set(self.rename_map.original[n] for n in cprior
                                      if n <= clust_start and self.orig.has_node(n))
                    if cluster_set:
                        clusters.append(cluster_set)
            finally:
                nx.relabel_nodes(self.dendrogram, self.rename_map.original, copy=False)
        return sorted(clusters, key=lambda c: -len(c))

def grid_2d(dim):
    """Creates a 2d grid of dimension dim"""
    graph = nx.grid_2d_graph(dim, dim)

    for node in graph:
        graph.node[node]['asn'] = 1
        graph.node[node]['x'] = node[0] * 150
        graph.node[node]['y'] = node[1] * 150
        graph.node[node]['device_type'] = 'router'
        graph.node[node]['platform'] = 'cisco'
        graph.node[node]['syntax'] = 'ios_xr'
        graph.node[node]['host'] = 'internal'
        graph.node[node]['ibgp_role'] = "Peer"

    mapping = {node: "%s_%s" % (node[0], node[1]) for node in graph}
    # Networkx wipes data if remap with same labels
    nx.relabel_nodes(graph, mapping, copy=False)
    for src, dst in graph.edges():
        graph[src][dst]['type'] = "physical"
        # add global index for sorting

    SETTINGS['General']['deploy'] = True
    SETTINGS['Deploy Hosts']['internal'] = {
        'cisco': {
            'deploy': True,
        },
    }

    return graph

    def _get_sys_graph(self):
        """Return the subsystem graph for this Group."""

        sgraph = self._relevance._sgraph
        if self.pathname:
            path = self.pathname.split('.')
            start = self.pathname + '.'
        else:
            path = []
            start = ''
        graph = sgraph.subgraph([n for n in sgraph
                                  if n.startswith(start)])
        renames = {}
        for node in graph.nodes_iter():
            renames[node] = '.'.join(node.split('.')[:len(path)+1])
            if renames[node] == node:
                del renames[node]

        # get the graph of direct children of current group
        nx.relabel_nodes(graph, renames, copy=False)

        # remove self loops created by renaming
        graph.remove_edges_from([(u, v) for u, v in graph.edges()
                                 if u == v])
        return graph

def load_graphs(fname, atlas):
    """Load a graph and return both the original and the common-nodes one.

    This reads a pickled graph and returns a pair: the graph just loaded and
    one that contains only nodes common with those in the atlas.
    """
    f = open(fname, 'rb')
    f.seek(0)                   # make sure we're at the beginning of the file
    g_coco = pickle.load(f)
    f.close()

    # relabel nodes to remove 'PHT00-' prefix
    remove_pht_map = dict([(lab, lab.replace('PHT00-','')) for lab in g_coco])
    g_coco = nx.relabel_nodes(g_coco, remove_pht_map)

    if hasattr(atlas, 'cocomac'):
        # relabel graph according to cocomac nodes
        name_map = dict(zip(atlas.cocomac, atlas.label))
        g = nx.relabel_nodes(g_coco, name_map)
    else:
        g = g_coco
        
    common = set(g.nodes()).intersection(set(atlas.label))

    gnorm = nx.DiGraph()
    for node in common:
        gnorm.add_node(node)
    for u,v,data in g.edges(data=True):
        if u in common and v in common:
                gnorm.add_edge(u, v, data)

    return g, gnorm

    def get_clusters(self, num_clusters=None):
        if num_clusters == None:
            index, value = max(enumerate(self.quality_history), key=lambda iv: iv[1])
            num_clusters = len(self.quality_history) - index

        nx.relabel_nodes(self.dendrogram, self.rename_map.integer, copy=False)
        start_node = max(self.dendrogram)

        priors, fringe = self.dendrogram_crawl(start=start_node, max_steps=num_clusters - 1)

        # Double check we got the right number of values
        if len(fringe) != num_clusters:
            raise ValueError(
                "get_clusters failed to retrieve "
                + "%d clusters correctly (got %d instead)" % (num_clusters, len(fringe))
            )

        clusters = []
        for neg_clust_start in fringe:
            clust_start = -neg_clust_start
            cprior, cfringe = self.dendrogram_crawl(start=clust_start, priors=priors.copy())
            clusters.append(
                set(self.rename_map.original[n] for n in cprior if n <= clust_start and self.orig.has_node(n))
            )
        nx.relabel_nodes(self.dendrogram, self.rename_map.original, copy=False)
        return sorted(clusters, key=lambda c: -len(c))

def coword_network(mesh_df, start, end,topic_count=0):
        """
        constructs a coword network for the years supplied;
        nodes will be labelled by topic, have a 'weight' of co-occurrence,
        a 'start_year' attribute,
        and an 'end_year' attribute which is the end year of the search

        Parameters
        ----------------
        mesh_df: a dataframe with at least the topics and years columns
        start: start year
        end: end year
        topic_count: the number of the topics to use
        (not too big, otherwise coword matrix will be huge
        """

        # determine the number of topics to count
        all_topics = [t for top in mesh_df.topics.dropna() for t in top]
        topic_collection = collections.Counter(all_topics)
        if topic_count > 0 and topic_count < len(topic_collection):
            common_topics = [k[0] for k in topic_collection.most_common(topic_count)]
        else:
            common_topics = sorted(topic_collection.keys())

        cow_df = coword_matrix_years(mesh_df, start, end, common_topics)
        cow_nx = nx.from_numpy_matrix(cow_df.as_matrix())
        col_names = cow_df.columns.tolist()
        labels = {col_names.index(l): l for l in col_names}
        start_year = {i: end for i in range(0, len(col_names))}
        end_year = {i: start for i in range(0, len(col_names))}
        nx.set_node_attributes(cow_nx, 'start_year', start_year)
        nx.set_node_attributes(cow_nx, 'end_year', end_year)
        nx.relabel_nodes(cow_nx, labels, copy=False)
        return cow_nx

def grid_2d(dim):
    import networkx as nx
    graph = nx.grid_2d_graph(dim, dim)

    for n in graph:
        graph.node[n]['asn'] = 1
        graph.node[n]['x'] = n[0] * 150
        graph.node[n]['y'] = n[1] * 150
        graph.node[n]['device_type'] = 'router'
        graph.node[n]['platform'] = 'cisco'
        graph.node[n]['syntax'] = 'ios_xr'
        graph.node[n]['host'] = 'internal'
        graph.node[n]['ibgp_level'] = 0

    mapping = {n: "%s_%s" % (n[0], n[1]) for n in graph}
    nx.relabel_nodes(graph, mapping, copy=False) # Networkx wipes data if remap with same labels
    for index, (src, dst) in enumerate(graph.edges()):
        graph[src][dst]['type'] = "physical"
        graph[src][dst]['edge_id'] = "%s_%s_%s" % (index, src, dst) # add global index for sorting

    SETTINGS['General']['deploy'] = True
    SETTINGS['Deploy Hosts']['internal'] = {
        'cisco': {
        'deploy': True,
        },
    }

    return graph

def grid_2d(dim):
    """Creates a 2d grid of dimension dim"""
    graph = nx.grid_2d_graph(dim, dim)

    for node in graph:
        graph.node[node]["asn"] = 1
        graph.node[node]["x"] = node[0] * 150
        graph.node[node]["y"] = node[1] * 150
        graph.node[node]["device_type"] = "router"
        graph.node[node]["platform"] = "cisco"
        graph.node[node]["syntax"] = "ios_xr"
        graph.node[node]["host"] = "internal"
        graph.node[node]["ibgp_role"] = "Peer"

    mapping = {node: "%s_%s" % (node[0], node[1]) for node in graph}
    # Networkx wipes data if remap with same labels
    nx.relabel_nodes(graph, mapping, copy=False)
    for src, dst in graph.edges():
        graph[src][dst]["type"] = "physical"
        # add global index for sorting

    SETTINGS["General"]["deploy"] = True
    SETTINGS["Deploy Hosts"]["internal"] = {"cisco": {"deploy": True}}

    return graph

def synthetic_three_level(n,p1,p2,p3,J_isolates=False,F_isolates=False,D_isolates=False):#,isolate_up=True,isolate_down=True):
    
    k=n

    J=nx.erdos_renyi_graph(n,p1) #The first layer graph
    Jis = nx.isolates(J)
    F=nx.erdos_renyi_graph(n,p2) #The second layer graph
    Fis = nx.isolates(F)
    D=nx.erdos_renyi_graph(n,p3) #The third layer graph
    Dis = nx.isolates(D)

    def translation_graph(J,F,D):
        H1=nx.Graph()
        H2=nx.Graph()
        for i in range(n):
            H1.add_edges_from([(J.nodes()[i],F.nodes()[i])])
            H2.add_edges_from([(F.nodes()[i],D.nodes()[i])])
        return H1, H2

    Jed = set(J.edges())
    Fed = set(F.edges())
    Ded = set(D.edges())
    l=[Jed,Fed,Ded]
    lu = list(set.union(*l))
    JFD=nx.Graph()
    JFD.add_edges_from(lu)

    G=nx.Graph()  #The synthetic two-layer graph
    
    # Relabing nodes maps
    
    mappingF={}
    for i in range(2*n):
        mappingF[i]=n+i
    FF=nx.relabel_nodes(F,mappingF,copy=True)
    
    mappingD={}
    for i in range(2*n):
        if i >n-1:
            mappingD[i]=i-n
        else:
            mappingD[i]=2*n+i
    DD=nx.relabel_nodes(D,mappingD,copy=True)
    
    H1, HH2 = translation_graph(J,FF,DD)
    
    G.add_edges_from(J.edges())
    G.add_edges_from(H1.edges())
    G.add_edges_from(DD.edges())
    G.add_edges_from(HH2.edges())
    G.add_edges_from(FF.edges())

    edgeList = []
    for e in H1.edges():
        edgeList.append(e)
    for e in HH2.edges():
        edgeList.append(e)
    
    return G, J, FF, DD, JFD, edgeList  

 def as_nx_graph(self):
     result = nx.from_numpy_matrix(self.matrix)
     nx.relabel_nodes(result,
                      dict(zip(
                          numpy.arange(len(result.nodes())),
                          self.rows)),
                      False)
     return result

    def _setup_graphs(self, group, connections):
        """
        Set up dependency graphs for variables and components in the Problem.

        Returns
        -------
        tuple of (nx.DiGraph, nx.DiGraph)
            The variable graph and the component graph.

        """
        params_dict = self.params_dict
        unknowns_dict = self.unknowns_dict

        vgraph = nx.DiGraph()  # var graph
        sgraph = nx.DiGraph()  # subsystem graph

        compins = {}  # maps input vars to components
        compouts = {} # maps output vars to components

        promote_map = {}

        # ensure we have system graph nodes even for unconnected subsystems
        sgraph.add_nodes_from([s.pathname for s in group.subsystems(recurse=True)])

        for meta in params_dict.values():
            param = meta['pathname']
            tcomp = param.rsplit('.', 1)[0]
            compins.setdefault(tcomp, []).append(param)
            if param in connections and meta['promoted_name'] != param:
                promote_map[param] = meta['promoted_name']

        for meta in unknowns_dict.values():
            unknown = meta['pathname']
            scomp = unknown.rsplit('.', 1)[0]
            compouts.setdefault(scomp, []).append(unknown)
            if meta['promoted_name'] != unknown:
                promote_map[unknown] = meta['promoted_name']

        for target, source in connections.items():
            vgraph.add_edge(source, target)
            sgraph.add_edge(source.rsplit('.', 1)[0], target.rsplit('.', 1)[0])

        # connect inputs to outputs on same component in order to fully
        # connect the variable graph.
        for comp, inputs in compins.items():
            for inp in inputs:
                for out in compouts.get(comp, ()):
                    vgraph.add_edge(inp, out)

        # now collapse any var nodes with implicit connections
        nx.relabel_nodes(vgraph, promote_map, copy=False)

        # remove any self edges created by the relabeling
        for u, v in vgraph.edges():
            if u == v:
                vgraph.remove_edge(u, v)

        return vgraph, sgraph