import sys
import os

"""
Usage:

$ python dfa_dot.py <path_to_log_file> <output_folder>

e.g.

$ cd protocol_dfa
$ python scripts/dfa_dot.py results/test_dfa_dot/extracted_dfa.log results/test_dfa_dot  skip_state

"""

inp_file = sys.argv[1]
outfolder = sys.argv[2]
skip_state = int(sys.argv[3])

outname = ".".join(inp_file.split(".")[:-1]).split("/")[-1]

dot_file = os.path.join(outfolder,  outname + '.dot')
img_file = os.path.join(outfolder,  outname + '.png')

if not os.path.exists(outfolder):
	os.makedirs(outfolder)

class Transition(object):
	def __init__(self, c_state, inp, n_state):
		self.c_state = c_state
		self.inp = inp
		self.n_state = n_state


transitions = []
states = []
with open(inp_file, 'r+') as ifile:
	count = 0
	for line in ifile:
		line = line[:-1]
		line = line.split(' ')
		print(line)
		if count == 0:
			if line[0] != 'x':
				raise 'The input file must start with \'x\''
			states = [ int(j) for j in line[2:]]
		else:
			inp = int(line[0])
			for idx in range(len(states)):
				c_state = states[idx]
				n_state = int(line[idx+2])
				##Remove this condition
				if c_state != skip_state and n_state != skip_state:
					transitions.append(Transition(c_state, inp, n_state))
				  			
		count += 1

print("writing dot file at ", dot_file)
with open(dot_file, 'w+') as ofile:
	ofile.write("digraph dfa {\n")
	for tr in transitions:
		ofile.write("S{} -> S{} [label=\"{}\"];\n".format( str(tr.c_state),str(tr.n_state), str(tr.inp) ))
	ofile.write("}\n")

print("creating image at ", img_file)
cmd = "dot -Tpng {} > {}".format(dot_file, img_file)
os.system(cmd)