#epidemics.py
import networkx as nx
import random

class Model_ep:
	def __init__(self, dyngraph, infected):
		self.G = dyngraph
		self.I = infected
		self.S = []
		self.R = []
		self.E = []
		self.beta = 0.5
		self.gamma = 0.5
		# self.model = model
		self.states = {}
		self.nodes = []
		self.nodestate = zip(self.nodes, self.states)
		self.I_period = 100
		self.E_period = 50
		self.summary = []


	def model_config(self):
		for g in self.G:
			for node in g.nodes():
				self.nodes.append(node)
		self.nodes = list(self.nodes)
		# print(self.nodes)
		for node in self.nodes:
			if node in self.I.keys():
				# self.states.append(1)
				pass
			else:
				# self.states.append(0)
				self.S.append(node)
				# print(self.S)
		# self.nodestate = zip(nodes, states)
	

	def random_toss(self):
		random.seed()
		r = random.random()
		if r < self.beta:
			return True
		else:
			return False


	def simulate_SIR(self):
		self.model_config()
		self.summary.append((len(self.S), len(self.I), len(self.R)))
		counter = 0
		for g in self.G:
			for node in g.nodes():
				if node in self.I.keys():
					for n in g.neighbors(node):
						if n not in self.I.keys():
							# print(n)
							# print(self.I.keys())
							# print(self.S)
							r = self.random_toss()
							if r:
								self.I.update({n: 0})
								self.S.remove(n)
		for node in self.I.keys():
			if self.I[node] == self.I_period:
				self.I.update({node: 0}) 
				self.R.append(node)
			elif self.I[node] != 0:
				self.I.update({node: self.I[node] + 1})
		return self.summary









			


