import spacy
from spacy.tokens import Doc
import logging, sys, time
from lib.CoNLL_Annotation import CoNLL09_Token
import my_utils.file_utils as fu
from germalemma import GermaLemma


TIGER_CORPUS = "/home/daza/datasets/TIGER_conll/tiger_release_aug07.corrected.16012013.conll09"
OUTPUT_FILE = "/home/daza/datasets/TIGER_conll/tiger_spacy_parsed.conllu"
TEXT_OUTPUT = "/home/daza/datasets/TIGER_conll/tiger_all.txt"


class WhitespaceTokenizer(object):
	def __init__(self, vocab):
		self.vocab = vocab

	def __call__(self, text):
		words = text.split(' ')
		# All tokens 'own' a subsequent space character in this tokenizer
		spaces = [True] * len(words)
		return Doc(self.vocab, words=words, spaces=spaces)


def get_conll_str(spacy_doc):
	conll_lines = [] # We want: [ID, FORM, LEMMA, UPOS, XPOS, FEATS, HEAD, DEPREL, DEPS, MISC]
	for ix, token in enumerate(spacy_doc):
		content = (str(ix), token.text, find_germalemma(token.text, token.tag_, token.lemma_), token.pos_, token.tag_, "_", "_", "_", "_", "_")
		conll_lines.append("\t".join(content))
	return "\n".join(conll_lines)


# def freeling_lemma_lookup():
# 	dicts_path = "/home/daza/Frameworks/FreeLing/data/de/dictionary/entries/"
	
def find_germalemma(word, pos, spacy_lemma):
	simplify_pos = {"ADJA":"ADJ", "ADJD":"ADJ",
					"NA":"N", "NE":"N", "NN":"N",
					"ADV":"ADV", "PAV":"ADV", "PROAV":"ADV", "PAVREL":"ADV", "PWAV":"ADV", "PWAVREL":"ADV",
					"VAFIN":"V", "VAIMP":"V", "VAINF":"V", "VAPP":"V", "VMFIN":"V", "VMINF":"V",
					"VMPP":"V", "VVFIN":"V", "VVIMP":"V", "VVINF":"V", "VVIZU":"V","VVPP":"V"
				}
	# simplify_pos = {"VERB": "V", "ADV": "ADV", "ADJ": "ADJ", "NOUN":"N", "PROPN": "N"}
	try:
		return lemmatizer.find_lemma(word, simplify_pos.get(pos, "UNK"))
	except:
		return spacy_lemma


if __name__ == "__main__":
	file_has_next, chunk_ix = True, 0
	CHUNK_SIZE = 10000
	
	# =====================================================================================
	#                    LOGGING INFO ...
	# =====================================================================================
	logger = logging.getLogger(__name__)
	console_hdlr = logging.StreamHandler(sys.stdout)
	file_hdlr = logging.FileHandler(filename=f"logs/Parse_Tiger.SpaCy.log")
	logging.basicConfig(level=logging.INFO, handlers=[console_hdlr, file_hdlr])
	logger.info(f"Chunking TIGER Corpus in chunks of {CHUNK_SIZE} Sentences")
	
	# =====================================================================================
	#                    POS TAG DOCUMENTS
	# =====================================================================================
	spacy_de = spacy.load("de_core_news_lg", disable=["ner", "parser"])
	spacy_de.tokenizer = WhitespaceTokenizer(spacy_de.vocab) # We won't re-tokenize to respect how the source CoNLL are tokenized!
	write_out = open(OUTPUT_FILE, "w")
	lemmatizer = GermaLemma()
	write_plain = open(TEXT_OUTPUT, "w")
	
	start = time.time()
	total_processed_sents = 0
	line_generator = fu.file_generator(TIGER_CORPUS)
	while file_has_next:
		sents, gld, file_has_next = fu.get_file_text_chunk(line_generator, chunk_size=CHUNK_SIZE, token_class=CoNLL09_Token)
		total_processed_sents += len(sents)
		logger.info(f"Already processed {total_processed_sents} sentences...")
		for doc in spacy_de.pipe(sents, batch_size=1000, n_process=10):
			conll_str = get_conll_str(doc)
			write_out.write(conll_str)
			write_out.write("\n\n")
			write_plain.write(" ".join([x.text for x in doc])+"\n") # OPTIONAL: This can be commented for efficiency...
			
	end = time.time()
	logger.info(f"Processing File {TIGER_CORPUS} took {(end - start)} seconds!")