import youGen_NG

# define logging types
gen_tree = 1
parameter_trees = 2

class Gentree:
	'''
	Design decisions
		count_stack is used to keep track of the counts rooted at each
		 rule.
	'''
	def __init__(self,file_name):
		'''
		Purpose
			create data structures for calculating the counts and
			 writing generation tree to file.
		Precondition
			file_name is valid.
			must be called before language generation.
		'''
		self.count_stack = []
		self.file_handle = open(file_name,'w')

	def destructor(self):
		'''
		Purpose
			destroy data structures for calculating the counts and
			 writing generation tree to file.
		Precondition
			must be called after language generation.
		'''
		self.count_stack = None
		self.file_handle.close()

	def get_count(self,i):
		'''
		Purpose
			return the counts rooted at rule of index i in the rule
			 stack.
		Precondition
			i is between 0 and size of rule stack-1
		'''
		if i < len(self.count_stack):
			return self.count_stack[i]
		# log_start is not called for next rule yet
		return 0

	def log_start(self,logging_type,R,S,is_ground):
		'''
		Purpose
			log the start of the application of rule R.
		Precondition
			logging_types is gen_tree or parameter_trees.
			R is a Rule object.
			S is a Sentential_form object.
			S is the sentential form just after the application of
			 rule R.
			must be called after application of rule R.
			is_ground is either True or False.
		'''
		global gen_tree,parameter_trees

		self.count_stack.append(0)

		if is_ground:
			# increments the counts
			for i in range(len(self.count_stack)):
				self.count_stack[i] += 1

		if logging_type == gen_tree:
			self.file_handle.write('<gentree>\n')
			self.file_handle.write('<id>' + str(R) + '</id>\n')
			self.file_handle.write('<s>' + str(S.contents) + \
			 '</s>\n')
		elif logging_type == parameter_trees:
			self.file_handle.write('<parameter_trees>\n')

	def log_end(self,logging_type):
		'''
		Purpose
			log the end of the application of rule R.
		Precondition
			logging_types is gen_tree or parameter_trees.
			must be called before application of alternative rule.
		'''
		global gen_tree,parameter_trees

		count = self.count_stack.pop()

		if logging_type == gen_tree:
			self.file_handle.write('<count>' + str(count) + \
			 '</count>\n')
			self.file_handle.write('</gentree>\n')
		elif logging_type == parameter_trees:
			self.file_handle.write('</parameter_trees>\n')
			# want number of gentree instead of parameter values
			for i in range(len(self.count_stack)):
				self.count_stack[i] -= count
