
__author__ = "Aric Hagberg (hagberg@lanl.gov)"
#    Copyright (C) 2004-2008 by 
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

__all__ = ['write_dot', 'read_dot', 'graphviz_layout', 'pydot_layout',
           'to_pydot', 'from_pydot']

import sys
from networkx.utils import _get_fh
import networkx


def write_dot(G,path):
    
    try:
        import pydot
    except ImportError:
        raise ImportError, \
          "write_dot() requires pydot http://dkbza.org/pydot.html/"
    fh=_get_fh(path,'w')
    P=to_pydot(G)
    fh.write(P.to_string())
    fh.flush() # might be a user filehandle so leave open (but flush)
    return

def read_dot(path):
    
    try:
        import pydot
    except ImportError:
        raise ImportError, \
          "read_dot() requires pydot http://dkbza.org/pydot.html/"

    fh=_get_fh(path,'r')
    data=fh.read()        
    P=pydot.graph_from_dot_data(data)
    return from_pydot(P)


def from_pydot(P):
    

    if P.get_strict(None): # pydot bug: get_strict() shouldn't take argument 
        multiedges=False
    else:
        multiedges=True
        
    if P.get_type()=='graph': # undirected
        if multiedges:
            create_using=networkx.MultiGraph()
        else:
            create_using=networkx.Graph()
    else:
        if multiedges:
            create_using=networkx.MultiDiGraph()
        else:
            create_using=networkx.DiGraph()

    # assign defaults        
    N=networkx.empty_graph(0,create_using)
    N.name=P.get_name()
    node_attr={}

    # add nodes, attributes to N.node_attr
    for p in P.get_node_list():
        n=p.get_name()
        if n.startswith('"'):
            n=n[1:-1]
        if n in ('node','graph','edge'):
            continue
        N.add_node(n,**p.get_attributes())

    # add edges
    for e in P.get_edge_list():
        u=e.get_source()
        v=e.get_destination()
        attr=e.get_attributes()
        N.add_edge(u,v,**attr)

    # add default attributes for graph, nodes, edges
    N.graph['graph']=P.get_attributes()
    # get atributes not working for these?
    # get_node_defaults()
    N.graph['node']={}
    if 'node' in P.obj_dict['nodes']:
        N.graph['node']=P.obj_dict['nodes']['node'][0]['attributes']
    # get_edge_defaults()
    N.graph['edge']={}
    if 'edge' in P.obj_dict['nodes']:
        N.graph['edge']=P.obj_dict['nodes']['edge'][0]['attributes']
    N.node_attr=node_attr

    return N        

def to_pydot(N, strict=True):
    
    try:
        import pydot
    except ImportError:
        raise ImportError, \
          "to_pydot() requires pydot http://dkbza.org/pydot.html/"

    # set Graphviz graph type
    if N.is_directed():
        graph_type='digraph'
    else:
        graph_type='graph'
    strict=N.number_of_selfloops()==0 and not N.is_multigraph() 

    P = pydot.Dot(graph_type=graph_type,strict=strict)

    for n,nodedata in N.nodes_iter(data=True):
        str_nodedata=dict((k,str(v)) for k,v in nodedata.iteritems())
        p=pydot.Node(str(n),**str_nodedata)
        P.add_node(p)

    if N.is_multigraph():
        for u,v,key,edgedata in N.edges_iter(data=True,keys=True):
            str_edgedata=dict((k,str(v)) for k,v in edgedata.iteritems())
            edge=pydot.Edge(str(u),str(v),key=str(key),**str_edgedata)
            P.add_edge(edge)
        
    else:
        for u,v,edgedata in N.edges_iter(data=True):
            str_edgedata=dict((k,str(v)) for k,v in edgedata.iteritems())
            edge=pydot.Edge(str(u),str(v),**str_edgedata)
            P.add_edge(edge)

    try:
        P.obj_dict['attributes'].update(N.graph.get('graph',{}))
    except:
        pass
    try:
        P.obj_dict['nodes']['node'][0]['attributes'].update(N.graph.get('node',{}))
    except:
        pass
    try:
        P.obj_dict['nodes']['edge'][0]['attributes'].update(N.graph.get('edge',{}))
    except:
        pass

    return P


def pydot_from_networkx(N):
	
        from warnings import warn
        warn('pydot_from_networkx is replaced by to_pydot', DeprecationWarning)
        return to_pydot(N)

def networkx_from_pydot(D, create_using=None):
	
        from warnings import warn
        warn('networkx_from_pydot is replaced by from_pydot', 
             DeprecationWarning)
        return from_pydot(D)

def graphviz_layout(G,prog='neato',root=None, **kwds):
    
    return pydot_layout(G=G,prog=prog,root=root,**kwds)


def pydot_layout(G,prog='neato',root=None, **kwds):
    
    try:
        import pydot
    except ImportError:
        raise ImportError, \
          "pydot_layout() requires pydot http://dkbza.org/pydot.html/"

    P=to_pydot(G)
    if root is not None :
        P.set("root",str(root))

    D=P.create_dot(prog=prog)

    if D=="":  # no data returned
        print "Graphviz layout with %s failed"%(prog)
        print
        print "To debug what happened try:"
        print "P=pydot_from_networkx(G)"
        print "P.write_dot(\"file.dot\")"
        print "And then run %s on file.dot"%(prog)
        return

    Q=pydot.graph_from_dot_data(D)

    node_pos={}
    for n in G.nodes():
        node=Q.get_node(pydot.Node(n).get_name())
        pos=node.get_pos()[1:-1] # strip leading and trailing double quotes
        if pos != None:
            xx,yy=pos.split(",")
            node_pos[n]=(float(xx),float(yy))
    return node_pos

