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

G = nx.read_edgelist("out.moreno_zebra_zebra", comments="%")
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()

# simple example for diffusion
def print_G(G, S):
	colors = []
	for v in G.nodes:
		if S[v] == 1:
			colors.append('green')
		elif S[v] == 2:
			colors.append('red')
		else:
			colors.append('blue')
	
	nx.draw_kamada_kawai(G, with_labels=True, node_color=colors)
	plt.show()

def init_state(G):
	S = {}
	for v in G.nodes():
		S[v] = 0
		
	# for i in range(0,1):
	# 	S[random.choice(list(G.nodes()))] = 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,1):
		S[C[i][0]] = 1
		
	for i in range(0,1):
		u = random.choice(list(G.nodes()))
		while S[u] == 1:
			u = random.choice(list(G.nodes()))
		S[u] = 2
		
	return S

def run_defusion(G):
	S = init_state(G)
	print_G(G, S)
	
	changes = 1
	while changes > 0:
		changes = 0
		for v in G.nodes():
			ones = 0
			twos = 0
			for u in G.neighbors(v):
				if S[u] == 1:
					ones += 1
				elif S[u] == 2:
					twos += 1
			
			orig = S[v]
			if ones > twos:
				S[v] = 1
			elif twos > ones:
				S[v] = 2
			elif ones > 0 and twos > 0:
				if random.random() < 0.5:
					S[v] = 1
				else:
					S[v] = 2
				
			if orig != S[v]:
				changes += 1
				
		print(changes)
		print_G(G, S)

run_defusion(G)


# 0 = healthy
# 1 = infected
# 2 = dead
# 3 = survived/immune

# t = infection duration
# p = transmission rate
# d = death rate upon infection

def init_state_epidemic(G):
	S = {}
	for v in G.nodes():
		S[v] = (0,0)
		
	#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]] = (3,0)
		
	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)
	
	# u = random.choice(list(G.nodes()))
	# while S[u][0] == 3:
	# 	u = random.choice(list(G.nodes()))
	# S[u] = (1,1)
	
	return S


def run_epidemic(G, p, d, t):
	S = init_state_epidemic(G)
	
	go = 1
	infections = 1
	deaths = 0
	duration = 0
	while go > 0:
		go = 0
		duration += 1
		for v in G.nodes():
			if S[v][1] == 0:
				continue
				
			for u in G.neighbors(v):
				if S[u][0] == 0 and random.random() < p:
					S[u] = (1,1)
					infections += 1
					
			S[v] = (S[v][0], S[v][1]+1)
			if S[v][1] == (t + 1):
				if random.random() < d:
					S[v] = (2,0)
					deaths += 1
				else:
					S[v] = (3,0)
			else:
				go += 1
				
		print("day:", duration, "; infections: ", infections)
		
	return (infections, deaths, duration)


G = nx.read_edgelist("Slashdot0902.txt")
G = G.subgraph(sorted(nx.connected_components(G), key=len, reverse=True)[0])

p = 0.01
d = 0.03
t = 14
i_count = 0
d_count = 0
t_count = 0
iterations = 1
for i in range(0, iterations):
	(infections, deaths, duration) = run_epidemic(G, p, d, t)
	i_count += infections
	d_count += deaths
	t_count += duration

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