#!/usr/bin/env python

from igraph import Graph


def unify_vertex_ids(g1, g2, attr):
    """Given two graphs g1 and g2, extends them with isolated vertices and
    permutes the vertices of g2 to ensure that both graphs have the same
    vertex ID set in the same order, assuming that the vertex attribute
    named ``attr`` stores the vertex IDs.
    """
    ids1, ids2 = g1.vs[attr], g2.vs[attr]
    assert len(set(ids1)) == len(ids1), "Vertex IDs in g1 must be unique"
    assert len(set(ids2)) == len(ids2), "Vertex IDs in g2 must be unique"
    ids_missing_from_g1 = sorted(set(ids2) - set(ids1))
    ids_missing_from_g2 = sorted(set(ids1) - set(ids2))
    if ids_missing_from_g1:
        g1 = g1.copy()
        n1 = g1.vcount()
        g1.add_vertices(len(ids_missing_from_g1))
        g1.vs[n1:][attr] = ids_missing_from_g1
    if ids_missing_from_g2:
        g2 = g2.copy()
        n2 = g2.vcount()
        g2.add_vertices(len(ids_missing_from_g2))
        g2.vs[n2:][attr] = ids_missing_from_g2
    reverse_index = dict((v, k) for k, v in enumerate(g1.vs[attr]))
    permutation_vector = [reverse_index[id] for id in g2.vs[attr]]
    g2 = g2.permute_vertices(permutation_vector)
    return g1, g2


def union_by_attr(g1, g2, attr):
    """Takes the union of two graphs g1 and g2 based on the given vertex
    attribute.
    """
    g1, g2 = unify_vertex_ids(g1, g2, attr)
    result = g1.union(g2)
    result.vs[attr] = g1.vs[attr]
    return result


g1 = Graph.GRG(100, 0.2)
g1.vs["id"] = [i * 3 for i in xrange(100)]

g2 = Graph.GRG(100, 0.2)
g2.vs["id"] = [i * 5 for i in xrange(100)]

merged_graph = union_by_attr(g1, g2, "id")
print(merged_graph.vs["id"])