#!/usr/bin/python3 

import networkx as nx
import matplotlib.pyplot as plt
from collections import OrderedDict
from operator import getitem
import itertools, os

# AGraph preserves attributes, networkx Graph does not.
# Many of the desired functions are in networkx.
# So import AGraph to keep attributes, then convert to Networkx.
A = nx.drawing.nx_agraph.to_agraph(nx.drawing.nx_pydot.read_dot("./1_mo_color_DOTFILE.dot"))
A.layout('dot')
A.draw('tree.png')
A.remove_node('\\n') # Remove "newline" node from newline end of dot file
G=nx.DiGraph(A)

color_map = []
color_d = {}
node_pos = {}

# Compartments
S = 0
I = 0
E = 0
R = 0
D = 0

# Params
beta = 0
delta = 0
gamma_r = 0
gamma_d = 0
mu = 0
epsilon = 0
omega = 0

for node in A:
    color = A.get_node(node).attr.to_dict()['fillcolor']
    str_pos = A.get_node(node).attr.to_dict()['pos']
    coords = str_pos.split(',')
    x = coords[0] # layout for draw function
    y = coords[1]
    node_pos[node] = float(x), float(y)
    if color is None or color == '':
        color_map.append("white")
        color_d[node] = color
        in_edges = list(G.in_edges(node))
    elif color == 'yellow':
        color_map.append(color)
        color_d[node] = color
        in_edges = list(G.in_edges(node))
    else:
        color_map.append(color)
        color_d[node] = color
        in_edges = list(G.in_edges(node))