# Source code for prov.graph

    from __future__ import (absolute_import, division, print_function,
                            unicode_literals)

    __author__ = 'Trung Dong Huynh'
    __email__ = 'trungdong@donggiang.com'

    import networkx as nx

    from prov.model import *

    INFERRED_ELEMENT_CLASS = {
        PROV_ATTR_ENTITY: ProvEntity,
        PROV_ATTR_ACTIVITY: ProvActivity,
        PROV_ATTR_AGENT: ProvAgent,
        PROV_ATTR_TRIGGER: ProvEntity,
        PROV_ATTR_GENERATED_ENTITY: ProvEntity,
        PROV_ATTR_USED_ENTITY: ProvEntity,
        PROV_ATTR_DELEGATE: ProvAgent,
        PROV_ATTR_RESPONSIBLE: ProvAgent,
        PROV_ATTR_SPECIFIC_ENTITY: ProvEntity,
        PROV_ATTR_GENERAL_ENTITY: ProvEntity,
        PROV_ATTR_ALTERNATE1: ProvEntity,
        PROV_ATTR_ALTERNATE2: ProvEntity,
        PROV_ATTR_COLLECTION: ProvEntity,
        PROV_ATTR_INFORMED: ProvActivity,
        PROV_ATTR_INFORMANT: ProvActivity
    }

    def prov_to_graph(prov_document):
        """ Convert a :class:`~prov.model.ProvDocument` to a `MultiDiGraph
        <http://networkx.github.io/documentation/latest/reference/classes.multidigraph.html>`_
        instance of the `NetworkX <https://networkx.github.io/>`_ library.

        :param prov_document: The :class:`~prov.model.ProvDocument` instance to convert.
        """
        g = nx.MultiDiGraph()
        unified = prov_document.unified()
        node_map = dict()
        for element in unified.get_records(ProvElement):
            g.add_node(element)
            node_map[element.identifier] = element

        for relation in unified.get_records(ProvRelation):
            # taking the first two elements of a relation
            attr_pair_1, attr_pair_2 = relation.formal_attributes[:2]
            # only need the QualifiedName (i.e. the value of the attribute)
            qn1, qn2 = attr_pair_1[1], attr_pair_2[1]
            if qn1 and qn2:  # only proceed if both ends of the relation exist
                try:
                    if qn1 not in node_map:
                        node_map[qn1] = \
                            INFERRED_ELEMENT_CLASS[attr_pair_1[0]](None, qn1)
                    if qn2 not in node_map:
                        node_map[qn2] = \
                            INFERRED_ELEMENT_CLASS[attr_pair_2[0]](None, qn2)
                except KeyError:
                    # Unsupported attribute; cannot infer the type of the element
                    continue  # skipping this relation
                g.add_edge(node_map[qn1], node_map[qn2], relation=relation)
        return g

    def graph_to_prov(g):
        """ Convert a `MultiDiGraph <http://networkx.github.io/documentation/latest/reference/classes.multidigraph.html>`_
            back to a :class:`~prov.model.ProvDocument`.

            :param g: The graph instance to convert.
        """
        prov_doc = ProvDocument()
        for n in g.nodes_iter():
            if isinstance(n, ProvRecord) and n.bundle is not None:
                prov_doc.add_record(n)
        for _, _, edge_data in g.edges_iter(data=True):
            try:
                relation = edge_data['relation']
                if isinstance(relation, ProvRecord):
                    prov_doc.add_record(relation)
            except KeyError:
                pass

        return prov_doc