#!/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())

def prep_seirds(weighted):
  print("Prepping the SEIRDS model, using trivial weighting=", weighted)
  # 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
  inf_ct = 0 # infection rate counter
  recov_ct = 0 # recov rate counter
  
  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))
          out_edges = list(G.out_edges(node))
          
          tmp_S = 1
          tmp_recov = 0
          tmp_inf = 0
          for source in in_edges:
              tmp_S = 1
              # If previous node was infected, then we are recovered
              if (color_d[source[0]] == 'red'):
                if(weighted == 'False' or not in_edges):
                  tmp_recov = 1
                else:
                  recov_ct = recov_ct + 1/len(in_edges) # trivial weighting
          recov_ct = recov_ct + tmp_recov
          for source in out_edges:
            if (color_d[source[0]] == 'red'):
             if(weighted == 'False' or not out_edges):
                tmp_inf = 1
            else:
                inf_ct = inf_ct + 1/len(out_edges) # trivial weighting
          inf_ct = inf_ct + tmp_inf
          S = S + tmp_S
      elif color == 'yellow':
          color_map.append(color)
          color_d[node] = color
          in_edges = list(G.in_edges(node))
          
          tmp_E = 1
          tmp_R = 0
          tmp_inf = 0
          tmp_recov = 0
          for source in in_edges:
              # If previous node was infected, then we are recovered
              if (color_d[source[0]] == 'red'):
                  tmp_E = 0
                  if(weighted == 'False' or not in_edges):
                    tmp_recov = 1
                  else:
                    recov_ct = recov_ct + 1/len(in_edges) # trivial weighting
              if (color_d[source[0]] == '' or color_d[source[0]] == 'white'):
                  if(weighted == 'False' or not in_edges):
                    tmp_inf = 1 # add 1 for the inf counter
                  else:
                    inf_ct = inf_ct + 1/len(in_edges) # trivial weighting
          E = E + tmp_E
          recov_ct = recov_ct + tmp_recov
          inf_ct = inf_ct + tmp_inf
      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
          else:
            tmp_inf = 0
            for source in in_edges:
              if (color_d[source[0]] == '' or color_d[source[0]] == 'white'):
                if(weighted == 'False' or not in_edges):
                  tmp_inf = 1
                else:
                  inf_ct = inf_ct + 1/len(in_edges) # trivial weightin
            inf_ct = inf_ct + tmp_inf
  
  # Params
  beta = (inf_ct)/len(A) # rate of infec
  delta = 1 # incubation period
  #gamma_r = R/len(A) # recov rate
  gamma_r = recov_ct/len(A)
  gamma_d = D/len(A) # death rate
  mu = D/(I_R+I_D) # fatality ratio
  epsilon = ep_tmp/len(A) # infected import rate
  omega = 1 # 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("infect counter:", str(inf_ct))
  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))
  print("omega:", str(omega))
  
  return (S, E, I_R+I_D, R, D, beta, delta, gamma_r, gamma_d, mu, epsilon, omega)