import codecs
import networkx as nx
import pydot
import argparse

def create_graph(sim,k):
	G = nx.DiGraph()
	'''with open(rlabel,'') as r:
		for line in r:
			#G.add_node('"'+line.strip().split(':')[-1]+'"')
			G.add_node('"'+line.strip().replace(':','%')+'"')
			'''
	cur = ''
	count =0
	
	with open(sim,'r') as s:
		for line in s:
			if not line.startswith('\t'):
				#cur = '"'+line.strip().split(':')[-1]+'"'
				cur = '"'+line.strip().replace(':','%')+'"'
				count = 0
			else:
				if k > count:
					l = line.strip().split(';')
					#if not G.has_edge(cur,'"'+l[0].split(':')[-1]+'"'):
					#if not G.has_edge(cur,'"'+l[0].replace(':','%')+'"'):
						#G.add_edge(cur,'"'+l[0].split(':')[-1]+'"',weight=l[1])
					G.add_edge(cur,'"'+l[0].replace(':','%')+'"',weight=l[1])
				count += 1
	
	return G

def get_subgraph(G,node,dist):
#	dic = nx.single_source_shortest_path_length(G,'"'+node.replace(':','%')+'"',dist)

#	nodes = dic.keys()
	G_UND = G.to_undirected()
	dic = nx.single_source_shortest_path_length(G_UND,'"'+node.replace(':','%')+'"',dist)
	nodes = dic.keys()

	sub = G.subgraph(nodes)
	return sub

#write to file, read from file

def draw(G,file):
	p = nx.to_pydot(G)
	p.write_jpg(file)

def save(G,file):
	nx.write_gpickle(G,file)

def read(file):
	return nx.read_gpickle(file)

if __name__ == '__main__':
		parser = argparse.ArgumentParser()
		parser.add_argument('--sim-file', dest = 'sim', required=True,help = 'File with labeled similarity results')
		parser.add_argument('--out', dest = 'out', required=True, help = 'Outfile')
		parser.add_argument('-k', dest = 'k', required=True, help = 'k most similar words')
		parser.add_argument('-w', '--word', dest = 'word', required=True, help = 'central word')
		parser.add_argument('-n', dest = 'n', required=True, help = 'depth of graph')

		args = parser.parse_args()

		G = create_graph(args.sim,int(args.k))

		#save(G,'tymcz')
		#G = read('tymcz')

		H = get_subgraph(G,args.word,int(args.n))
		draw(H,args.out)

