#!/usr/bin/python3 

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

# Change dir to location of this python file
print(os.getcwd())
#os.chdir(os.path.dirname(sys.argv[0]))
#print(os.getcwd())

# 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 = {} # used for drawing/graphing

# Compartments
S = 0
I_R = 0
I_D = 0
E = 0
R = 0
D = 0

ep_tmp = 0 # counter for epsilon

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))
        tmp_S = 1
        for source in in_edges:
            tmp_S = 1
            # If previous node was infected, then we are recovered
            if (color_d[source[0]] == 'red'):
                R = R + 1
                tmp_S = 0
                break # No need to check the other nodes
        S = S + tmp_S
            #G[source[0]][node]['weight'] = 3 
    elif color == 'yellow':
        color_map.append(color)
        color_d[node] = color
        in_edges = list(G.in_edges(node))
        tmp_E = 1
        for source in in_edges:
            tmp_E = 1
            # If previous node was infected, then we are recovered
            if (color_d[source[0]] == 'red'):
                R = R + 1
                tmp_E = 0
                break # No need to check the other nodes
        E = E + tmp_E
    else:
        color_map.append(color)
        color_d[node] = color
        # Check if node dies
        out_edges = list(G.out_edges(node))
        if not out_edges:
            D = D + 1
            I_D = I_D + 1
        else:
            I_R = I_R + 1
        # Check if imported
        in_edges = list(G.in_edges(node))
        if not in_edges:
            ep_tmp = ep_tmp + 1

# Params
beta = (I_R+I_D)/len(A) # rate of infec (I/total?)
delta = E/len(A) # symptom appearance rate (E/total?)
gamma_r = R/len(A) # recov rate (R/total?)
gamma_d = D/len(A) # death rate (D/total?)
mu = D/(I_R+I_D) # fatality ratio (D/I)
epsilon = ep_tmp/len(A) # infected import rate
omega = 0 # waning immunity rate


print("Model Compartments:")
print("S:", str(S))
print("E:", str(E))
print("I_R:", str(I_R))
print("I_D:", str(I_D))
print("R:", str(R))
print("D:", str(D))
print("\n")

print("Model Parameters:")
print("beta:", str(beta))
print("delta:", str(delta))
print("gamma_r:", str(gamma_r))
print("gamma_d:", str(gamma_d))
print("mu:", str(mu))
print("epsilon:", str(epsilon))