from .NounChunk import NounChunk
from .Entity import Entity
from .BasicToken import BasicToken
from .EntityChunk import EntityChunk
from .EntityType import EntityType
from .remove_list_duplicates import remove_list_duplicates
from .join_punctuation import join_punctuation

from vineyard import Graph, NodeStyle
import spacy

nlp = spacy.load('en_core_web_sm')

class Token(BasicToken):
	def __add__(self, other):
		"""
		:param Token or Document other: adding two tokens (or several tokens) creates a Document
		:rtype: Document
		"""
		return Document(text=str(self).rstrip()+' '+str(other).lstrip())

	@property
	def parents(self):
		"""
		:rtype: list[Token]
		"""
		try:
			return self._parents
		except AttributeError:
			self._parents = [Token(obj=x, document=self.document) for x in self.token.ancestors]
			return self._parents

	@property
	def children(self):
		"""
		:rtype: list[Token]
		"""
		try:
			return self._children
		except AttributeError:
			self._children = [Token(obj=x, document=self.document) for x in self.token.children]
			return self._children



class Document:
	def __init__(self, text, id=0):
		"""
		:type text: str
		"""
		if isinstance(text, list):
			text = ' '.join([str(x) for x in text])
		self._id = id
		self._text = str(text)
		self._tokens = remove_list_duplicates([Token(obj=x, document=self) for x in self.doc])
		self._entities = remove_list_duplicates([Entity(obj=span, document=self) for span in self.doc.ents])
		self._noun_chunks = remove_list_duplicates([
			NounChunk(obj=span, document=self) for span in self.doc.noun_chunks
		])
		self._entity_chunks = remove_list_duplicates([EntityChunk(entity=entity) for entity in self.entities])

	@property
	def id(self):
		return self._id

	@property
	def doc(self):
		"""
		:rtype: tokens.Doc
		"""
		try:
			return self._doc
		except AttributeError:
			self._doc = nlp(self._text)
			return self._doc

	def __repr__(self):
		return self.doc.__repr__()

	def __str__(self):
		return self.doc.__str__()

	def graph_str(self):
		return f'd: {self}'[0:50]


	@property
	def tokens(self):
		"""
		:rtype: list[Token]
		"""
		return self._tokens

	@property
	def entities(self):
		"""
		:rtype: list[Entity]
		"""
		return self._entities

	@property
	def noun_chunks(self):
		"""
		:rtype: list[NounChunk]
		"""
		return self._noun_chunks

	@property
	def entity_chunks(self):
		"""
		:rtype: list[EntityChunk]
		"""
		return self._entity_chunks

	@property
	def entity_graph(self):
		"""
		:rtype: Graph
		"""
		try:
			return self._entity_graph
		except AttributeError:
			graph = Graph(ordering=False)

			# add tokens to the graph
			for token in self.tokens:
				graph.add_node(
					name=str(token.id), label=f'{token.graph_str()}',
					style=NodeStyle(text_size=8, shape='rect', style='"rounded, filled"')
				)

			# add noun_chunks to the graph
			for noun_chunk in self.noun_chunks:
				graph.add_node(
					name=str(noun_chunk.id), label=f'{noun_chunk.graph_str()}',
					style=NodeStyle(
						fill_colour='lightpink', text_colour='black', text_size=8,
						shape='rect', style='"rounded, filled"'
					)
				)

				# noun_chunk to its tokens
				for token in noun_chunk.tokens:
					graph.connect(start=str(token.noun_chunk.id), end=str(token.id))

			# add entities to the graph
			for entity in self.entities:
				graph.add_node(
					name=str(entity.id), label=f'{entity.graph_str()}',
					style=NodeStyle(fill_colour='gold3', text_colour='black')
				)
				# connect entity to its tokens
				for token in entity.tokens:
					graph.connect(start=str(token.id), end=str(entity.id))

			# add entity_chunks to the graph
			for entity_chunk in self.entity_chunks:
				graph.add_node(
					name=str(entity_chunk.id), label=f'{entity_chunk.graph_str()}',
					style=NodeStyle(fill_colour='gold', text_colour='black')
				)

				# connect entity_chunk to its entities
				for entity in entity_chunk.entities:
					graph.connect(start=str(entity.id), end=str(entity_chunk.id))

			self._entity_graph = graph
			return self._entity_graph

	@property
	def graph(self):
		"""
		:rtype: Graph
		"""
		try:
			return self._graph
		except AttributeError:
			graph = Graph(ordering=False)
			for token in self.tokens:
				graph.add_node(name=str(token.id), label=token.graph_str())
			for token in self.tokens:
				for child in token.children:
					dependency_label = str(child.dependency).replace('_', '\n')
					graph.connect(start=str(token.id), end=str(child.id), label=dependency_label)
			self._graph = graph
			return self._graph

	def __add__(self, other):
		"""
		:type other: Document or Token
		:rtype: Document
		"""
		return self.__class__(text=str(self).rstrip()+' '+str(other).lstrip())

	def bind(self, chunk_type='entity_chunk'):
		"""
		:param str chunk_type: can be one of entity_chunk, noun_chunk, or entity
		:rtype: list[Token or EntityChunk or NounChunk or Entity]
		"""
		if chunk_type=='entity_chunk':
			chunks = [x for x in self.tokens if x.entity_chunk is None] + self.entity_chunks
		elif chunk_type=='entity':
			chunks = [x for x in self.tokens if x.entity is None] + self.entities
		elif chunk_type=='noun_chunk':
			chunks = [x for x in self.tokens if x.noun_chunk is None] + self.noun_chunks
		else:
			raise ValueError(f'chunk_type "{chunk_type}" is not acceptable!')
		return sorted(chunks, key=lambda x: x.start)

	def mask_entity_types(self, entity_types, chunk_type='entity_chunk'):
		"""
		:param list[str] or list[EntityType] or str entity_types: entity types to be masked, it can be 'all'
		:param str chunk_type: can be one of entity_chunk, noun_chunk, or entity
		:rtype: str
		"""
		if entity_types=='all':
			pass
		else:
			entity_types = [EntityType(name=x) if not isinstance(x, EntityType) else x for x in entity_types]
		chunks = self.bind(chunk_type=chunk_type)
		strings = []
		for chunk in chunks:
			string = str(chunk)
			if 	(chunk_type=='entity_chunk' and isinstance(chunk, EntityChunk)) or \
				(chunk_type=='noun_chunk' and isinstance(chunk, NounChunk)):
				for entity_type in sorted(chunk.entity_types):
					if entity_types=='all':
						string = f'<{entity_type.long_name.upper()}>'
						break
					elif entity_type in entity_types:
						string = f'<{entity_type.long_name.upper()}>'
						break
			elif chunk_type=='entity' and isinstance(chunk, Entity):
				if entity_types=='all':
					string = f'<{chunk.entity_type.long_name.upper()}>'
				elif chunk.entity_type in entity_types:
					string = f'<{chunk.entity_type.long_name.upper()}>'
			strings.append(string)
		return ' '.join(join_punctuation(seq=[str(x) for x in strings]))



