import networkx as nx
import matplotlib.pyplot as plt
import random
import operator

################################################################################
# Like we did before in this context, let's look at our Zebra interaction graph
G = nx.read_edgelist("out.moreno_zebra_zebra.data", comments="%")

# We'll grab the largest connected component out of it
# We assume that diffusive process can't jump without an edge
G = G.subgraph(sorted(nx.connected_components(G), key=len, reverse=True)[0])

# Look at centrality versus the visual network topology
# nx.degree_centrality(G)
nx.closeness_centrality(G)
# nx.betweenness_centrality(G)
# nx.eigenvector_centrality(G)

nx.draw(G, with_labels=True)
plt.show()


################################################################################
# We'll consider a basic SIR simulation below
#
# Vertex state possibilities:
# 0 = susceptible
# 1 = infected
# 2 = removed
#
# Model parameters:
# t = infection duration
# p = transmission rate

# We'll initialize state based off of centrality measures
# States are stored as a (state, duration) tuple
def init_state_epidemic(G):
	S = {}
	for v in G.nodes():
		S[v] = (0,0)
		
	# We can look at a couple things:
	# 1. How does centrality of initial infected impact the epidemic?
	# 2. How does centrality of initial removed impact the epidemic?
	
	for i in range(0,10):
		S[random.choice(list(G.nodes()))] = (1,1)
	#
	#C = nx.degree_centrality(G)
	#C = nx.betweenness_centrality(G)
	#C = nx.closeness_centrality(G)
	#C = nx.eigenvector_centrality(G)
	#
	# C = sorted(C.items(), key=operator.itemgetter(1), reverse=True)
	# for i in range(0,10):
	# 	S[C[i][0]] = (1,1)
	#
	# C = sorted(C.items(), key=operator.itemgetter(1), reverse=True)
	# for i in range(0,10):
	# 	S[C[i][0]] = (2,0)
	
	return S

# Here we'll run the epidemic by updating states of each vertex across some 
# number of iterations. We're done when there are no longer any vertices with 
# the state of infected.
def run_epidemic(G, p, t):
	S = init_state_epidemic(G)
	infected = 1
	infections = 1
	duration = 0
	while infected > 0:
		infected = 0
		duration += 1
		for v in G.nodes():
			
			# We'll only consider infected vertices
			if S[v][1] == 0:
				continue
			
			# Infect our susceptible neighbors with probability p
			for u in G.neighbors(v):
				if S[u][0] == 0 and random.random() < p:
					S[u] = (1,1)
					infections += 1
			
			# We remove ourself if t is greater than the infectiousness duration
			# Otherwise, we need to keep processing an additional iteration
			S[v] = (S[v][0], S[v][1]+1)
			if S[v][1] == (t + 1):
					S[v] = (2,0)
			else:
				infected += 1
		print("Day:", duration, "; Infectious: ", infected, "; Total infected: ", infections)
	return (infections, duration)

################################################################################
# We'll consider a couple interaction graph for our diffusive process
# G = nx.read_edgelist("Slashdot0902.data")
# G = G.subgraph(sorted(nx.connected_components(G), key=len, reverse=True)[0])
G = nx.read_weighted_edgelist("out.dnc-temporalGraph.data", create_using=nx.MultiGraph(), comments="%")
G = nx.Graph(G) # get rid of multi-edges

# We'll use probability of infection to be 1% chance per interaction, with
# duration of infection 10 days. For simplicity, we can assume that p does not
# change relative to how long one is infected, and one is infectious the entire
# time they are infect as well. Also, we'll consider the graph we create to be
# defining daily interactions -- not going to consider a temporal graph.
# What we want to look at:
# - How many total people end up infected based on who we initialize as infected
# - How long does the disease last until there are no more infections
p = 0.01
t = 10
i_count = 0
d_count = 0
t_count = 0
iterations = 1
for i in range(0, iterations):
	(infections, duration) = run_epidemic(G, p, t)
	i_count += infections
	t_count += duration

i_count /= iterations
t_count /= iterations
print("Average infections:", i_count)
print("Average duration:", t_count)