#!/usr/bin/python

# This program file used to generate initial probability distribution of each individual in population, and HCU's for their location.
# Also it generates transition probability matrices for ovarall population and HCUs.

import igraph
from igraph import *
import pickle
import os
import numpy as np, numpy.random
import sys
from random import shuffle

numpy.set_printoptions(threshold=sys.maxint)
dir_path = os.path.dirname(os.path.realpath(__file__))
#global variables
num_nodes    = 50           # no. of vertices in graph
num_edges    = 1225         # no. of edges in graph
num_HCUs     = 10            # HCU = Health Care Units
num_people   = 20000        # total population over all cities
num_infected = 400	    # number of people infected initially

resources_path = dir_path +"/../Input/ESP_Resources/resources/"
BA_path = dir_path +"/../Input/ESP_Input/BA/"

#fractions = [0.2, 0.2, 0.2, 0.2, 0.2]
#fractions = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]

fractions = [0.10269572, 0.091075994, 0.06969163, 0.057720345, 0.046038687, 0.03835277, 0.0371144, 0.036875892, 0.025788365, 0.02514214, 0.023251563, 0.022824375, 0.019855658, 0.016803905, 0.016182477, 0.015012438, 0.014841966, 0.014273625, 0.013892641, 0.01375648, 0.013503628, 0.013320478, 0.013217802, 0.013087938, 0.012885103, 0.0122730415, 0.011593596, 0.010804294, 0.010622481, 0.010287265, 0.0100797005, 0.00991943, 0.009844957, 0.009667816, 0.009587185, 0.009349478, 0.009239837, 0.009220168, 0.008859861, 0.008849304, 0.008702181, 0.008695323, 0.008537289, 0.0085336575, 0.008336963, 0.0082676895, 0.007951862, 0.007930055, 0.00785025, 0.00779032]

# format the data generated by random_graph_generation function
def format_string3(gen_data, start_index):
	#format_data = str(gen_data).strip('[ ').strip(']').replace('\n','')
	format_data = "";
	for x in gen_data:
		num = int(x) + start_index
		format_data = format_data + str(num) + ", "
	format_data = format_data[:-2]
        return format_data

# format the data generated by random_graph_generation function
def format_string1(gen_data):
	format_data = str(gen_data).strip('[ ').strip(']').replace('\n','')
        return format_data

# format the data generated by generate_prob_distribution function
def format_string2(gen_data):
	format_data = str(gen_data).strip('[ ').strip(']').replace('\n','')
        temp_data = format_data.split()
        final_data = ""
        for x in temp_data:
		final_data = final_data + x + ", " 
	final_data = final_data.strip(', ')     
        return final_data
	
# this function generate initial probability distribution for individuals, and HCUs
def generate_prob_distribution(n):
	d = np.random.dirichlet(np.ones(n),size=1)
	return d

# this function generating the random graph for given numbe of vertices and edges
def random_graph_generation():
	random_graph = Graph.Erdos_Renyi(n=num_nodes, m=num_edges)
	#getting the adjacency matrix and adjacency list
	adj_matrix = random_graph.get_adjacency()
	adj_list = random_graph.get_adjlist()
        return (adj_matrix,adj_list)

# this function generating the random graph for given numbe of vertices and edges
def barabasi_graph_generation(_nodes):
	random_graph =  Graph.Barabasi(_nodes,2,outpref=False, directed=False, zero_appeal=1, implementation="psumtree", start_from=None)
	#getting the adjacency matrix and adjacency list
	adj_matrix = random_graph.get_adjacency()
	adj_list = random_graph.get_adjlist()
        return (adj_matrix,adj_list)


# this function write transition probability matrices of ovarall population, and HCUs into file
def write_trans_prob_matrix_of_population(adj_matrix, adj_list):
	# open file for writing for transition probability matrices of overall population based on random graph
        fp4 = open((resources_path+"transition_prob_matrix_of_people.txt"), "w+")
        #fp4.write("transition probability matrix of overall popuation : \n")
	fp4.write(str(num_people)+'\n')	
        row_id1 = 0
	for row_data1 in adj_matrix :
		row_len1  = len(adj_list[(row_id1)])
		temp_data = "" 	
                count1     = 0	
		for x1 in row_data1:
			if x1 == 1 :
				temp_data =  temp_data + str(0.3 / row_len1 ) + ', '
			else:
				if row_id1 == count1 :
					temp_data = temp_data + str(0.7) + ', '
				else:
					temp_data = temp_data + str(0.0) + ', '
			count1 = count1 + 1
                temp_data = temp_data.strip(', ')
		fp4.write(temp_data + "\n")
		row_id1 = row_id1 + 1	
	fp4.close()
   
# this function write transition probability matrices of ovarall population, and HCUs into file
def write_trans_prob_matrix_of_HCU(adj_matrix, adj_list):
	# open file for writing for transition probability matrices of HCUs based on random graph
        fp5 = open((resources_path+"transition_prob_matrix_of_HCUs.txt"), "w+")
	fp5.write(str(num_HCUs) + '\n')
        #fp5.write("transition probability matrix of all HCUs : \n")
        row_id2 = 0
	for row_data2 in adj_matrix :
		row_len2 = len(adj_list[(row_id2)])
		temp_data = "" 		
		for x2 in row_data2:
			if x2 == 1 :
				temp_data = temp_data + str(1.0 / row_len2 ) + ', '
			else:
				temp_data = temp_data + str(0.0) + ', '
                temp_data = temp_data.strip(', ')
		fp5.write(temp_data + "\n")
		row_id2 = row_id2 + 1	
	fp5.close()

# this function generate transition probability matrices for ovarall population, and HCUs
def barabasi_network_matrices():
	population = []
	total = 0
	for i in range(num_nodes):		
		if(i < 26): # if 2500->24, 10000->26, 20000->26
			num = int(num_people * fractions[i]) +1
		else :             
			num = int(num_people * fractions[i])  		 
		population.append(num)
 		print "\npeople at node %s : %s "%(i, population[i])
		#adj_matrix, adj_list = random_graph_generation()
		adj_matrix, adj_list = barabasi_graph_generation(num)
		#print adj_matrix
		#print adj_list
		filename = BA_path + "adjacency_list_of_node" + str(i) + ".txt"
        	# open file for writing adjacency matrix of random graph
        	fp3 = open(filename, "w+")
        	fp3.write(str(population[i])+ '\n')
		for row_data in adj_list :
			print format_string3(row_data, total)
			fp3.write(format_string3(row_data, total) + '\n')
		fp3.close()
		total = total + num
	print "\nTotal Population  : %s"%(total)
	fp3 = open((resources_path+"initial_population_at_each_node.txt"), "w+")
	fp3.write(str(total) + '\n')
	fp3.write(format_string1(population) + '\n')
	fp3.close()

              

# this function generate transition probability matrices for ovarall population, and HCUs
def transition_prob_matrices():
	adj_matrix, adj_list = random_graph_generation()
	#adj_matrix, adj_list = barabasi_graph_generation()
	print adj_matrix
	print adj_list
        # open file for writing adjacency matrix of random graph
        fp3 = open((resources_path + "adjacency_matrix_of_graph.txt"), "w+")
        fp3.write(str(num_nodes)+ '\n')
	for row_data in adj_matrix :
		print format_string1(row_data)
		fp3.write(format_string1(row_data) + '\n')
	fp3.close()
        write_trans_prob_matrix_of_population(adj_matrix, adj_list)
	write_trans_prob_matrix_of_HCU(adj_matrix, adj_list)

# this function generate initial probability distribution of each individual in population, HCU's for their location and list of people infected
def initial_prob_distribution(num_people, num_HCUs):
        # open file for writing initial probability distribution of each individual
	fp1 = open((resources_path+"initial_prob_dist_of_people.txt"), "w+") 	
	fp1.write(str(num_people)+'\n')	
	for i in range(num_people):
		prob_dist = generate_prob_distribution(num_nodes)
		fp1.write(format_string2(prob_dist) + '\n')
		#print prob_dist
		#print len(prob_dist[0])
	fp1.close()

	# open file for writing initial probability distribution of each HCU
	fp2 = open((resources_path +"initial_prob_dist_of_HCUs.txt"), "w+") 
	fp2.write(str(num_HCUs) + '\n')
	for i in range(num_HCUs):
		prob_dist = generate_prob_distribution(num_nodes)
		fp2.write(format_string2(prob_dist) + '\n')
		#print prob_dist
		#print len(prob_dist[0])
	fp2.close()	
	
	# open file for writing initial probability distribution of each HCU
	fp6 = open((resources_path+"initial_infected_people_list.txt"), "w+") 
	fp6.write(str(num_infected) + '\n')
	rs = np.random.choice(num_people, num_infected, replace=False)
	fp6.write(format_string2(rs) + '\n')		
	fp6.close()
         
def population_distribution_over_nodes():
	#open file for writing population distribution over nodes
	fp7 = open((resources_path+"initial_population_at_each_node.txt"), "w+") 
	fp7.write(str(num_people) + '\n')
	avg_pop   = num_people / num_nodes
	j = num_nodes/2
	extra_pop = (avg_pop/num_nodes)
	var = []

	if( num_nodes % 2 == 0):
 		for i in range(num_nodes):
			pop = avg_pop + ( extra_pop * (j-1) )
			var.append( int(pop))
			j = j-1	
		if( extra_pop != 0):
			var[(num_nodes-1)] = var[(num_nodes-1)] + avg_pop/2
	else:
		for i in range(num_nodes):
			pop = avg_pop + ( extra_pop * (j) )
			var.append( int(pop))
			j = j-1	
		#var[(num_nodes-1)] = var[(num_nodes-1)] + avg_pop/2

	shuffle(var)
	total = 0
	for i in range(num_nodes-1):
		fp7.write(str(var[i]) + ", ")
		total = total + var[i]
	total = total + var[(num_nodes-1)]
	fp7.write( str(var[(num_nodes-1)]) + "\n")
	fp7.close()	
	print total

# main function definition
def main():
	
	initial_prob_distribution(num_people, num_HCUs)
	transition_prob_matrices()	
	barabasi_network_matrices()
	#population_distribution_over_nodes()
# calling main function
main()

