import warnings
with warnings.catch_warnings():
	 warnings.simplefilter("ignore")
	 import networkit as nk
import networkx as nx

class DGraph(object):
	
	def __init__(self):
		
		self._GOR = None			# Object relation graph (edge = relation)
		self._GCR = None			# Class relation graph (edge = relation)
		self._GCM = None			# Class membership graph (edge = contains)
		
		self._GOR_names = {}		# {node_id: obj_id, ...}
		self._GOR_lookup = {}		# {obj_id: node_id, ...}
		self._GOR_data = {}  		# {obj_id: data, ...}
		self._GOR_labels = {}		# {(src_id, tgt_id): {label: weight, ...}, ...}
		
		self._GCR_names = {}		# {node_id: name, ...}
		self._GCR_lookup = {}		# {name: node_id, ...}
		self._GCR_data = {}  		# {name: data, ...}
		self._GCR_labels = {}		# {(src_id, tgt_id): [label, ...], ...}
		
		self._GCM_names = {}		# {node_id: name/obj_id, ...}
		self._GCM_lookup = {}		# {name/obj_id: node_id, ...}
		
		self.clear()
	
	def n_objects(self):
		
		return self._GOR.numberOfNodes()
	
	def n_classes(self):
		
		return self._GCR.numberOfNodes()
	
	def clear(self):
		
		self._GOR = nk.graph.Graph(directed = True)
		self._GCR = nk.graph.Graph(directed = True)
		self._GCM = nk.graph.Graph(directed = True)
		
		self._GOR_names.clear()
		self._GOR_lookup.clear()
		self._GOR_data.clear()
		self._GOR_labels.clear()
		
		self._GCR_names.clear()
		self._GCR_lookup.clear()
		self._GCR_data.clear()
		self._GCR_labels.clear()
		
		self._GCM_names.clear()
		self._GCM_lookup.clear()
	
	
	# ---- Objects
	# ------------------------------------------------------------------------
	
	def iter_objects(self):
		
		for obj_id in self._GOR_lookup:
			yield obj_id
	
	def iter_objects_data(self):
		
		for obj_id in self._GOR_data:
			yield self._GOR_data[obj_id]
	
	def iter_object_parents(self, obj_id):
		
		node_id = self._GCM_lookup[obj_id]
		for src in self._GCM.iterInNeighbors(node_id):
			yield self._GCM_names[src]
	
	def iter_object_ancestors(self, obj_id):
		
		node_id = self._GCM_lookup[obj_id]
		ancestors = set()
		queue = set([node_id])
		while queue:
			for src in self._GCM.iterInNeighbors(queue.pop()):
				if tgt not in ancestors:
					queue.add(tgt)
				ancestors.add(tgt)
		for node in ancestors:
			yield self._GCM_names[node]
	
	def has_object(self, obj_id):
		
		return obj_id in self._GOR_lookup
	
	def add_object(self, obj_id, data):
		
		if obj_id not in self._GOR_lookup:
			node_id = self._GOR.addNode()
			self._GOR_lookup[obj_id] = node_id
			self._GOR_names[node_id] = obj_id
			node_id = self._GCM.addNode()
			self._GCM_lookup[obj_id] = node_id
			self._GCM_names[node_id] = obj_id
		self._GOR_data[obj_id] = data
	
	def get_object_data(self, obj_id):
		
		return self._GOR_data[obj_id]
	
	def iter_object_relations(self, obj_id = None):
		
		if obj_id is None:
			for src, tgt in self._GOR.iterEdges():
				for label in self._GOR_labels[(src, tgt)]:
					yield self._GOR_names[src], self._GOR_names[tgt], label
		else:
			node_id = self._GOR_lookup[obj_id]
			for tgt in self._GOR.iterNeighbors(node_id):
				for label in self._GOR_labels[(node_id, tgt)]:
					yield self._GOR_names[tgt], label
	
	def add_object_relation(self, src_id, tgt_id, label, weight = None):
		
		src_id = self._GOR_lookup[src_id]
		tgt_id = self._GOR_lookup[tgt_id]
		
		if not self._GOR.hasEdge(src_id, tgt_id):
			self._GOR.addEdge(src_id, tgt_id)
			self._GOR_labels[(src_id, tgt_id)] = {}
		self._GOR_labels[(src_id, tgt_id)][label] = weight
	
	def has_object_relation(self, src_id, tgt_id, label):
		
		src_id = self._GOR_lookup[src_id]
		tgt_id = self._GOR_lookup[tgt_id]
		
		if not self._GOR.hasEdge(src_id, tgt_id):
			return False
		return label in self._GOR_labels[(src_id, tgt_id)]
	
	def get_object_relation_weight(self, src_id, tgt_id, label):
		
		src_id = self._GOR_lookup[src_id]
		tgt_id = self._GOR_lookup[tgt_id]
		
		if not self._GOR.hasEdge(src_id, tgt_id):
			return None
		return self._GOR_labels[(src_id, tgt_id)].get(label, None)
	
	def del_object_relation(self, src_id, tgt_id, label):
		
		src_id = self._GOR_lookup[src_id]
		tgt_id = self._GOR_lookup[tgt_id]
		
		if not self._GOR.hasEdge(src_id, tgt_id):
			return
		if label in self._GOR_labels[(src_id, tgt_id)]:
			del self._GOR_labels[(src_id, tgt_id)][label]
		if not self._GOR_labels[(src_id, tgt_id)]:
			self._GOR.removeEdge(src_id, tgt_id)
			del self._GOR_labels[(src_id, tgt_id)]
	
	def del_object(self, obj_id):
		
		node_id = self._GOR_lookup[obj_id]
		del self._GOR_names[node_id]
		del self._GOR_lookup[obj_id]
		del self._GOR_data[obj_id]
		for tgt in self._GOR.iterNeighbors(node_id):
			del self._GOR_labels[(node_id, tgt)]
		for src in self._GOR.iterInNeighbors(node_id):
			del self._GOR_labels[(src, node_id)]
		self._GOR.removeNode(node_id)
		
		node_id = self._GCM_lookup[obj_id]
		self._GCM.removeNode(node_id)
		del self._GCM_names[node_id]
		del self._GCM_lookup[obj_id]
	
	
	# ---- Classes
	# ------------------------------------------------------------------------
	
	def iter_classes(self):
		
		for name in self._GCR_lookup:
			yield name
	
	def iter_classes_data(self):
		
		for name in self._GCR_data:
			yield self._GCR_data[name]
	
	def has_class(self, name):
		
		return name in self._GCR_lookup
	
	def add_class(self, name, data):
		
		if name not in self._GCR_lookup:
			node_id = self._GCR.addNode()
			self._GCR_lookup[name] = node_id
			self._GCR_names[node_id] = name
			node_id = self._GCM.addNode()
			self._GCM_lookup[name] = node_id
			self._GCM_names[node_id] = name
		self._GCR_data[name] = data
	
	def rename_class(self, name, new_name):
		
		node_id = self._GCR_lookup[name]
		data = self._GCR_data[name]
		del self._GCR_lookup[name]
		del self._GCR_data[name]
		self._GCR_names[node_id] = new_name
		self._GCR_lookup[new_name] = node_id
		self._GCR_data[new_name] = data
		
		node_id = self._GCM_lookup[name]
		del self._GCM_lookup[name]
		self._GCM_names[node_id] = new_name
		self._GCM_lookup[new_name] = node_id
	
	def get_class_data(self, name):
		
		return self._GCR_data[name]
	
	def iter_class_children(self, name):
		
		node_id = self._GCM_lookup[name]
		for tgt in self._GCM.iterNeighbors(node_id):
			yield self._GCM_names[tgt]
	
	def iter_class_descendants(self, name):
		
		node_id = self._GCM_lookup[name]
		descendants = set()
		nk.traversal.Traversal.DFSfrom(self._GCM, node_id, lambda node: descendants.add(node))
		descendants.remove(node_id)
		for node in descendants:
			yield self._GCM_names[node]
	
	def iter_class_parents(self, child):
		
		node_id = self._GCM_lookup[child]
		for src in self._GCM.iterInNeighbors(node_id):
			yield self._GCM_names[src]
	
	def iter_class_ancestors(self, child):
		
		node_id = self._GCM_lookup[child]
		ancestors = set()
		queue = set([node_id])
		while queue:
			for src in self._GCM.iterInNeighbors(queue.pop()):
				if tgt not in ancestors:
					queue.add(tgt)
				ancestors.add(tgt)
		for node in ancestors:
			yield self._GCM_names[node]
	
	def iter_class_relations(self, name = None):
		
		if name is None:
			for src, tgt in self._GCR.iterEdges():
				for label in self._GCR_labels[(src, tgt)]:
					yield self._GCR_names[src], self._GCR_names[tgt], label
		else:
			node_id = self._GCR_lookup[name]
			for tgt in self._GCR.iterNeighbors(node_id):
				for label in self._GCR_labels[(node_id, tgt)]:
					yield self._GCR_names[tgt], label
	
	def add_class_child(self, name, child):
		# child = obj_id or subclass_name
		
		src_id = self._GCM_lookup[name]
		tgt_id = self._GCM_lookup[child]
		if not self._GCM.hasEdge(src_id, tgt_id):
			self._GCM.addEdge(src_id, tgt_id)
	
	def has_class_child(self, name, child):
		# child = obj_id or subclass_name
		
		src_id = self._GCM_lookup[name]
		tgt_id = self._GCM_lookup[child]
		
		return self._GCM.hasEdge(src_id, tgt_id)
	
	def has_class_relation(self, src_name, tgt_name, label):
		
		src_id = self._GCR_lookup[src_name]
		tgt_id = self._GCR_lookup[tgt_name]
		
		if not self._GCR.hasEdge(src_id, tgt_id):
			return False
		
		return label in self._GCR_labels[(src_id, tgt_id)]
	
	def add_class_relation(self, src_name, tgt_name, label):
		
		src_id = self._GCR_lookup[src_name]
		tgt_id = self._GCR_lookup[tgt_name]
		
		if not self._GCR.hasEdge(src_id, tgt_id):
			self._GCR.addEdge(src_id, tgt_id)
			self._GCR_labels[(src_id, tgt_id)] = []
		if label not in self._GCR_labels[(src_id, tgt_id)]:
			self._GCR_labels[(src_id, tgt_id)].append(label)
	
	def del_class_relation(self, src_name, tgt_name, label):
		
		src_id = self._GCR_lookup[src_name]
		tgt_id = self._GCR_lookup[tgt_name]
		
		if not self._GCR.hasEdge(src_id, tgt_id):
			return
		if label in self._GCR_labels[(src_id, tgt_id)]:
			self._GCR_labels[(src_id, tgt_id)].remove(label)
		if not self._GCR_labels[(src_id, tgt_id)]:
			self._GCR.removeEdge(src_id, tgt_id)
			del self._GCR_labels[(src_id, tgt_id)]
	
	def del_class_child(self, name, child):
		# child = obj_id or subclass_name
		
		src_id = self._GCM_lookup[name]
		tgt_id = self._GCM_lookup[child]
		if self._GCM.hasEdge(src_id, tgt_id):
			self._GCM.removeEdge(src_id, tgt_id)
	
	def del_class(self, name):
		
		node_id = self._GCR_lookup[name]
		del self._GCR_names[node_id]
		del self._GCR_lookup[name]
		del self._GCR_data[name]
		for tgt in self._GCR.iterNeighbors(node_id):
			del self._GCR_labels[(node_id, tgt)]
		for src in self._GCR.iterInNeighbors(node_id):
			del self._GCR_labels[(src, node_id)]
		self._GCR.removeNode(node_id)
		
		node_id = self._GCM_lookup[name]
		self._GCM.removeNode(node_id)
		del self._GCM_names[node_id]
		del self._GCM_lookup[name]
	
	
	# ---- Serialize
	# ------------------------------------------------------------------------
	
	def objects_to_nx(self):
		
		G = nx.MultiDiGraph()
		for obj_id in self._GOR_data:
			G.add_node(obj_id, data = self._GOR_data[obj_id])
		for src_id, tgt_id in self._GOR_labels:
			for label in self._GOR_labels[(src_id, tgt_id)]:
				weight = self._GOR_labels[(src_id, tgt_id)][label]
				if weight is None:
					G.add_edge(self._GOR_names[src_id], self._GOR_names[tgt_id], label)
				else:
					G.add_edge(self._GOR_names[src_id], self._GOR_names[tgt_id], label, weight = weight)
		return G
	
	def classes_to_nx(self):
		
		G = nx.MultiDiGraph()
		for name in self._GCR_data:
			G.add_node(name, data = self._GCR_data[name])
		for src_id, tgt_id in self._GCR_labels:
			for label in self._GCR_labels[(src_id, tgt_id)]:
				G.add_edge(self._GCR_names[src_id], self._GCR_names[tgt_id], label)
		return G
	
	def members_to_nx(self):
		
		G = nx.DiGraph()
		for name in self._GCM_lookup:
			G.add_node(name)
		for src_id, tgt_id in self._GCM.iterEdges():
			G.add_edge(self._GCM_names[src_id], self._GCM_names[tgt_id])
		return G
	
	def objects_from_nx(self, G):
		
		self._GOR = nk.graph.Graph(directed = True)
		self._GOR_names = {}		# {node_id: obj_id, ...}
		self._GOR_lookup = {}		# {obj_id: node_id, ...}
		self._GOR_data = {}  		# {obj_id: data, ...}
		self._GOR_labels = {}		# {(src_id, tgt_id): [[label, weight], ...], ...}
		for obj_id, data in G.nodes(data = "data"):
			node_id = self._GOR.addNode()
			self._GOR_lookup[obj_id] = node_id
			self._GOR_names[node_id] = obj_id
			self._GOR_data[obj_id] = data
		for src_id, tgt_id, label in G.edges(keys = True):
			weight = G.edges[src_id, tgt_id, label].get("weight", None)
			src_id = self._GOR_lookup[src_id]
			tgt_id = self._GOR_lookup[tgt_id]
			if not self._GOR.hasEdge(src_id, tgt_id):
				self._GOR.addEdge(src_id, tgt_id)
				self._GOR_labels[(src_id, tgt_id)] = {}
			self._GOR_labels[(src_id, tgt_id)][label] = weight
		
	def classes_from_nx(self, G):
		
		self._GCR = nk.graph.Graph(directed = True)
		self._GCR_names = {}		# {node_id: name, ...}
		self._GCR_lookup = {}		# {name: node_id, ...}
		self._GCR_data = {}  		# {name: data, ...}
		self._GCR_labels = {}		# {(src_id, tgt_id): [[label, weight], ...], ...}
		for name, data in G.nodes(data = "data"):
			node_id = self._GCR.addNode()
			self._GCR_lookup[name] = node_id
			self._GCR_names[node_id] = name
			self._GCR_data[name] = data
		for src_id, tgt_id, label in G.edges(keys = True):
			src_id = self._GCR_lookup[src_id]
			tgt_id = self._GCR_lookup[tgt_id]
			if not self._GCR.hasEdge(src_id, tgt_id):
				self._GCR.addEdge(src_id, tgt_id)
				self._GCR_labels[(src_id, tgt_id)] = []
			self._GCR_labels[(src_id, tgt_id)].append(label)
	
	def members_from_nx(self, G):
		
		self._GCM = nk.graph.Graph(directed = True)
		self._GCM_names = {}		# {node_id: name/obj_id, ...}
		self._GCM_lookup = {}		# {name/obj_id: node_id, ...}
		for name in G.nodes():
			node_id = self._GCM.addNode()
			self._GCM_lookup[name] = node_id
			self._GCM_names[node_id] = name
		for src_id, tgt_id in G.edges():
			src_id = self._GCM_lookup[src_id]
			tgt_id = self._GCM_lookup[tgt_id]
			self._GCM.addEdge(src_id, tgt_id)
		
	
	def objects_to_json(self, graph_attrs):
		
		return nx.node_link_data(self.objects_to_nx(), attrs = graph_attrs)
	
	def classes_to_json(self, graph_attrs):
		
		return nx.node_link_data(self.classes_to_nx(), attrs = graph_attrs)
	
	def members_to_json(self, graph_attrs):
		
		return nx.node_link_data(self.members_to_nx(), attrs = graph_attrs)
	
	def objects_from_json(self, data, graph_attrs):
		
		G = nx.node_link_graph(data, directed = True, multigraph = True, attrs = graph_attrs)
		self.objects_from_nx(G)
	
	def classes_from_json(self, data, graph_attrs):
		
		G = nx.node_link_graph(data, directed = True, multigraph = True, attrs = graph_attrs)
		self.classes_from_nx(G)
	
	def members_from_json(self, data, graph_attrs):
		
		G = nx.node_link_graph(data, directed = True, multigraph = False, attrs = graph_attrs)
		self.members_from_nx(G)
	
	
	def objects_to_pickle(self):
		
		return [
			self._GOR,
			self._GOR_names,
			self._GOR_lookup,
			self._GOR_data,
			self._GOR_labels,
		]
	
	def classes_to_pickle(self):
		
		return [
			self._GCR,
			self._GCR_names,
			self._GCR_lookup,
			self._GCR_data,
			self._GCR_labels,		
		]
	
	def members_to_pickle(self):
		
		return [
			self._GCM,
			self._GCM_names,
			self._GCM_lookup,
		]
	
	def objects_from_pickle(self, data):
		
		[
			self._GOR,
			self._GOR_names,
			self._GOR_lookup,
			self._GOR_data,
			self._GOR_labels,
		] = data
	
	def classes_from_pickle(self, data):
		
		[
			self._GCR,
			self._GCR_names,
			self._GCR_lookup,
			self._GCR_data,
			self._GCR_labels,		
		] = data
	
	def members_from_pickle(self, data):
		
		[
			self._GCM,
			self._GCM_names,
			self._GCM_lookup,
		] = data

