import random
import math
import sys

DEFAULT_NUM_VERTICES = 40000
DEFAULT_PROB_CONNECTED = 0.005
DEFAULT_DIRECTED_EDGES = True
DEFAULT_OUT_FILENAME = 'graph'
DEFAULT_NUM_OUT_FILES = 1

def main():
	args = sys.argv
	if len(args) == 5:
		edges = create_graph(int(args[1]), float(args[2]))
		write_to_file(edges, args[3], int(args[4]))
	if len(args) == 6:
		edges = create_graph(int(args[1]), float(args[2]), bool(args[5]))
		write_to_file(edges, args[3], int(args[4]))
	elif len(args) > 1:
		print('Usage: generate-random-graph.py <num_vertices> <probability_vertices_connected> <out_filename (no extension)> <num_out_files> <optional: directed_edges>')
	else:
		write_to_file(create_graph())
	print('done')

def create_graph(num_vertices=DEFAULT_NUM_VERTICES, prob_connected=DEFAULT_PROB_CONNECTED, directed_edges=DEFAULT_DIRECTED_EDGES):
	edges = set()
	num_edges = math.floor(num_vertices * num_vertices * prob_connected)
	for _ in range(num_edges):
		src_vertex = random.randint(1, num_vertices)
		dest_vertex = random.randint(1, num_vertices)
		while src_vertex == dest_vertex or (src_vertex, dest_vertex) in edges:
			dest_vertex = random.randint(1, num_vertices)
		edges.add((str(src_vertex), str(dest_vertex)))
		if not directed_edges:
			edges.add((dest_vertex, src_vertex))
	return convert_to_dict(edges)

def convert_to_dict(edges):
	neighbor_map = {}
	for (a, b) in edges:
		if a not in neighbor_map.keys():
			neighbor_map[a] = [b]
		else:
			neighbor_map[a].append(b)
	return neighbor_map

def write_to_file(neighbor_map, filename=DEFAULT_OUT_FILENAME, num_out_files=DEFAULT_NUM_OUT_FILES):
	num_nodes_per_file = math.ceil(len(neighbor_map) / num_out_files)
	neighbor_map_as_list = list(neighbor_map.items())
	count = 0

	for file_num in range(num_out_files):
		with open(f'{filename}{file_num + 1}.csv', 'w') as f:
			for i in range(num_nodes_per_file):
				if count >= len(neighbor_map_as_list):
					break
				f.write(neighbor_entry_to_string(neighbor_map_as_list[count][0], neighbor_map_as_list[count][0]) + '\n')
				count += 1

def neighbor_entry_to_string(vertex, neighbors):
    output = vertex + '\t'
    for neighbor in neighbors:
        output += neighbor + ','
    return output[0 : len(output) - 1]

if __name__ == '__main__':
	main()