import jinja2
from collections import namedtuple
import modelflow.elementarynodes
import modelflow.graph
import subprocess
import os
Node = namedtuple("node", "id label color")
Edge = namedtuple("edge", "start end label")

_template = """
digraph model {
  {% for node in nodes %}
     {{ node.id }} [label="{{ node.label }}" color={{ node.color }} shape=box];\
  {% endfor %}
  {% for edge in edges %}
     {{ edge.start }} -> {{ edge.end }} [label="{{ edge.label }}"];\
  {% endfor %}
}
"""


def draw(graph: modelflow.graph.DynamicGraph, file=None):
    if not file:
        file = os.path.join(os.environ['HOME'], "graph.dot")
    out_degs = graph._graph.out_degree()
    in_degs = graph._graph.in_degree()
    nodes = [Node(id(node), str(node), "black" if (in_degs[node] and out_degs[node]) else "blue" if out_degs[node] else "red") for node in graph._graph.nodes()]
    edges = [Edge(id(edge[0]), id(edge[1]),  edge[2]['to_channel']) for edge in graph._graph.edges(data=True)]
    res = jinja2.Template(_template).render(nodes=nodes, edges=edges)
    with open(file, 'wt') as f:
        f.write(res)
    subprocess.call(['dot', '-Tpng', '-O', file])


if __name__ == "__main__":
    minus = modelflow.elementarynodes.Minus()
    div = modelflow.elementarynodes.Division()
    graph = modelflow.graph.DynamicGraph()
    graph.add_inputs(["one", "two", "three"])
    graph.connect_input_to_node("one", minus, "plus")
    graph.connect_input_to_node("two", minus, "minus")
    graph.connect_input_to_node("three", div, "denominator")
    graph.connect_node_to_node(minus, div, "enumerator")
    graph.set_output(div)
    draw(graph)