
import os
import copy

class DecisionDiagram:
	# any operation that modifies the DD should take care of the root list
	root = []
	next_available_id = 0
	node_dict = dict()
	sigdict = dict()
	var_dict = dict({None:None})

	def __init__(self):
		self.root = []
		self.next_available_id = 0
		self.sigdict = dict()
		self.node_dict = dict()
		var_dict = dict({None:None})
		

	def __copy__(self):
		dd = DecisionDiagram()
		dd.root = copy.copy(self.root)
		dd.next_available_id = self.next_available_id
		dd.node_dict = self.node_dict.copy()
		dd.var_dict = self.var_dict.copy()
		return dd

	def __repr__(self):
		return "root: "+str(self.root) + "\n"+ str(self.node_dict) + "\n" + str(self.var_dict)+ "\n"

	def getRoot(self):
		return self.root

	def getPseudoRoot(self):
		self.node_dict[self.next_available_id] = (None, tuple(self.getRoot()))
		return self.next_available_id

	def addRoot(self, ids):
		assert isinstance(ids, list)
		self.root.extend(ids)

	def setRoot(self, id):
		self.root = [id]

	def isSingleRoot(self):
		return len(self.root)==1

	def refreshRoot(self):
		assert self.isPseudo(self.node_dict[self.next_available_id])
		del self.root
		self.root = list(self.node_dict[self.next_available_id][1])

	def isTerminal(self, node):
		return not isinstance(node[1][0], int)

	def isPseudo(self, node):
		return not self.isTerminal(node) and node[0]==None

	def vars(self):
		self.clean()
		var_set = set()
		for node in self.node_dict.values():
			if not node[0] == None:
				var_set.add(node[0])
		return var_set

	def refreshSigDict(self):
		self.clean()
		self.sigdict.clear()
		for key in self.node_dict.keys():
			if self.node_dict[key] not in self.sigdict: 
				self.sigdict[self.node_dict[key]] = key
			else:
				if not key==self.sigdict[self.node_dict[key]]:
					print self
					print key, self.sigdict[self.node_dict[key]]
					assert 0==1

	def addNewNode(self, in_var_id, in_edge_list, bound):
		self.node_dict[self.next_available_id] = (in_var_id, in_edge_list)
		self.next_available_id += 1
		self.var_dict[in_var_id] = bound
		return self.next_available_id-1

	def node_ids(self):
		def traverse_node_ids(subroot, visited):
			visited.add(subroot)
			if self.isTerminal(self.node_dict[subroot]):
				yield subroot
			else:
				for child in self.node_dict[subroot][1]:
					if child not in visited:
						for nid in traverse_node_ids(child, visited):
							yield nid
				yield subroot
		for nid in traverse_node_ids(self.getPseudoRoot(), set()):
			yield nid

	def edges(self):
		for nid in self.node_ids():
			if not self.isTerminal(self.node_dict[nid]):
				edge_dict = dict()
				for val, cid in enumerate(self.node_dict[nid][1]):
					if cid not in edge_dict:
						edge_dict[cid] = [val]
					else:
						edge_dict[cid].append(val)
				for cid, val in edge_dict.iteritems():
					yield nid, cid, val 
		
	def evaluate(self, assignment):
		# self.refreshSigDict()
		self.clean()
		for nid in self.node_ids():
			if not self.isTerminal(self.node_dict[nid]):
				edge_list = list(self.node_dict[nid][1])
				for i in range(len(edge_list)):
					child = edge_list[i]
					if self.node_dict[child][0] in assignment:
						edge_list[i] = self.node_dict[child][1][assignment[self.node_dict[child][0]]]
				self.node_dict[nid] = (self.node_dict[nid][0], tuple(edge_list))
		self.refreshRoot()
			

	def traverseRF(self, subroot, visited):
		visited.add(subroot)
		if self.isTerminal(self.node_dict[subroot]):
			return
		else:
			for child in self.node_dict[subroot][1]:
				if child not in visited:
					self.traverseRF(child, visited)
					
	def clean(self):
		ref_set = set()
		self.traverseRF(self.getPseudoRoot(), ref_set)
		for key in self.node_dict.keys():
			if key not in ref_set:
				del self.node_dict[key]

	def merge(self, other):
		merge_dict = dict()
		for key in other.node_ids():
			node = other.node_dict[key]
			if other.isPseudo(node):
				continue

			if other.isTerminal(node):
				sig = node
			else:
				edge_list = []
				for child in node[1]:
					edge_list.append(merge_dict[child])  
					
				sig = (node[0], tuple(edge_list))

			if sig not in self.sigdict:
				id = self.addNewNode(sig[0], sig[1], other.var_dict[sig[0]])
				merge_dict[key] = id
				self.sigdict[sig] = id
			else:
				merge_dict[key] = self.sigdict[sig]

		for root in other.getRoot():
			self.addRoot([merge_dict[root]])
		self.refreshSigDict()

	def diff(self, other):
		self.clean()
		other.clean()
		diff_nodes = []
		merge_dict = dict()
		self.refreshSigDict()
		for key, node in other.node_dict.iteritems():
			if other.isPseudo(node):
				continue

			if other.isTerminal(node):
				sig = node
			else:
				edge_list = []
				for child in node[1]:
					edge_list.append(merge_dict[child])  
				sig = (node[0], tuple(edge_list))

			if sig not in self.sigdict:
				diff_nodes.append(key)
				# for i in other.node_dict[key][1]:
				# 	print i, merge_dict[i]
				# print self.node_dict[4]
				merge_dict[key] = None
				# return diff_nodes
			else:
				merge_dict[key] = self.sigdict[sig]
		return diff_nodes


	def isConst(self):
		assert self.isSingleRoot()
		root = self.root[0]
		return self.isTerminal(self.node_dict[root])

	def getVal(self, assignment):
		assert self.isSingleRoot()
		id = self.root[0]
		while not self.isTerminal(self.node_dict[id]):
			id = self.node_dict[id][1][assignment[self.node_dict[id][0]]]
		return self.node_dict[id][1][0]

	def countPath(self):
		paths = dict()
		for edge in self.edges():
			source = edge[0]
			sink = edge[1]
			if source not in paths:
				paths[source] = 0
			if sink not in paths:
				if self.isTerminal(self.node_dict[sink]):
					paths[sink] = 1
				else:
					paths[sink] = 0
			paths[source] += paths[sink]
		return paths[self.getPseudoRoot()]

	
	def visualize(self, model_name):
		tmpfile = open(model_name+".dot", "w")
		tmpfile.write("digraph G { \n")
		for nid in self.node_ids():
			node = self.node_dict[nid]
			if not self.isTerminal(node):
				tmpfile.write(str(nid)+"[label="+"\""+str(nid)+" "+str([node[0]])+"\"]"+";\n")
			else:
				tmpfile.write(str(nid)+"[label="+"\""+str(nid)+" "+str(node[1][0])+"\"]"+";\n")

		for edge in self.edges():
			tmpfile.write(str(edge[0])+"-> "+str(edge[1])+"[label="+"\""+str(edge[2])+"\"]"+";\n")
			
		tmpfile.write("}")
		tmpfile.close()

		os.system("rm -f "+model_name+".png")
		print "Saving Decision Diagram to "+model_name+".png ..."
		os.system("dot "+model_name+".dot -Tpng >> "+model_name+".png")
		# os.system("gnome-open "+model_name+".png")

	def attach_right(self, val):
		for key, node in self.node_dict.iteritems():
			if self.isTerminal(node):
				self.node_dict[key] = (node[0], (node[1][0] + val,))
	def attach_left(self, val):
		for key, node in self.node_dict.iteritems():
			if self.isTerminal(node):
				self.node_dict[key] = (node[0], (val + node[1][0],))


