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

def format_morphological_features(token):
	"""
	Extract and format morphological features from a spaCy token for CoNLL-U output.
	
	Args:
		token: spaCy token object
		
	Returns:
		str: Formatted morphological features string for CoNLL-U 5th column
			 Returns "_" if no features are available
	"""
	if not hasattr(token, 'morph') or not token.morph:
		return "_"
	
	morph_dict = token.morph.to_dict()
	if not morph_dict:
		return "_"
	
	# Format as CoNLL-U format: Feature=Value|Feature2=Value2
	features = []
	for feature, value in sorted(morph_dict.items()):
		features.append(f"{feature}={value}")
	
	return "|".join(features)


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(anno_obj, spacy_doc, use_germalemma):
	#  First lines are comments. (metadata)
	conll_lines = anno_obj.metadata # Then we want: [ID, FORM, LEMMA, UPOS, XPOS, FEATS, HEAD, DEPREL, DEPS, MISC]
	for ix, token in enumerate(spacy_doc):
		morph_features = format_morphological_features(token)
		if use_germalemma == "True":
			content = (str(ix), token.text, find_germalemma(token.text, token.tag_, token.lemma_), token.pos_, token.tag_, morph_features, "_", "_", "_", "_")
		else:
			content = (str(ix), token.text, token.lemma_, token.pos_, token.tag_, morph_features, "_", "_", "_", "_") # Pure SpaCy!
		conll_lines.append("\t".join(content))
	return "\n".join(conll_lines)

	
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__":
	"""
		EXAMPLE:
		--- TIGER Classic Orthography ---
			python systems/parse_spacy.py --corpus_name Tiger --gld_token_type CoNLL09_Token \
				-i /home/daza/datasets/TIGER_conll/tiger_release_aug07.corrected.16012013.conll09 \
				-o /home/daza/datasets/TIGER_conll/tiger_spacy_parsed.conllu \
				-t /home/daza/datasets/TIGER_conll/tiger_all.txt
			
			python systems/parse_spacy.py --corpus_name TigerOld_test \
			-i /home/daza/datasets/TIGER_conll/data_splits/test/Tiger.OldOrth.test.conll \
			-o /home/daza/datasets/TIGER_conll/tiger_spacy_parsed.test.conllu
		
		--- TIGER New Orthography ---
			python systems/parse_spacy.py --corpus_name TigerNew \
				-i /home/daza/datasets/TIGER_conll/Tiger.NewOrth.train.conll \
				-o /home/daza/datasets/TIGER_conll/Tiger.NewOrth.train.spacy_parsed.conllu \
				-t /home/daza/datasets/TIGER_conll/Tiger.NewOrth.train.txt
			
			python systems/parse_spacy.py --corpus_name TigerNew_test \
			-i /home/daza/datasets/TIGER_conll/data_splits/test/Tiger.NewOrth.test.conll \
			-o /home/daza/datasets/TIGER_conll/Tiger.NewOrth.test.spacy_parsed.conllu
		
		--- German GSD Universal Deps ---
			python systems/parse_spacy.py --corpus_name DE_GSD \
				-i /home/daza/datasets/ud-treebanks-v2.2/UD_German-GSD/de_gsd-ud-test.conllu \
				-o /home/daza/datasets/ud-treebanks-v2.2/UD_German-GSD/de_gsd-ud-test.parsed.germalemma.conllu \
				-t /home/daza/datasets/ud-treebanks-v2.2/UD_German-GSD/de_gsd-ud-test.txt
				
			
		--- Real Data TEST  ---
		time python systems/parse_spacy.py --corpus_name DeReKo_a00 --comment_str "#" \
			-i /export/netapp/kupietz/N-GRAMM-STUDIE/conllu/a00.conllu.gz \
			-o /export/netapp/kupietz/N-GRAMM-STUDIE/conllu/0_SpaCyParsed/a00.spacy.gl.conllu
	"""
	
	parser = argparse.ArgumentParser()
	parser.add_argument("-i", "--input_file", help="Input Corpus", required=True)
	parser.add_argument("-n", "--corpus_name", help="Corpus Name", default="Corpus")
	parser.add_argument("-o", "--output_file", help="File where the Predictions will be saved", required=True)
	parser.add_argument("-t", "--text_file", help="Output Plain Text File", default=None)
	parser.add_argument("-sm", "--spacy_model", help="Spacy model containing the pipeline to tag", default="de_core_news_lg")
	parser.add_argument("-gtt", "--gld_token_type", help="CoNLL Format of the Gold Data", default="CoNLLUP_Token")
	parser.add_argument("-ugl", "--use_germalemma", help="Use Germalemma lemmatizer on top of SpaCy", default="True")
	parser.add_argument("-c", "--comment_str", help="CoNLL Format of comentaries inside the file", default="#")
	args = parser.parse_args()
	
	file_has_next, chunk_ix = True, 0
	CHUNK_SIZE = 20000
	SPACY_BATCH = 2000
	SPACY_PROC = 10
	
	# =====================================================================================
	#                    LOGGING INFO ...
	# =====================================================================================
	logger = logging.getLogger(__name__)
	console_hdlr = logging.StreamHandler(sys.stdout)
	file_hdlr = logging.FileHandler(filename=f"logs/Parse_{args.corpus_name}.SpaCy.log")
	logging.basicConfig(level=logging.INFO, handlers=[console_hdlr, file_hdlr])
	logger.info(f"Chunking {args.corpus_name} Corpus in chunks of {CHUNK_SIZE} Sentences")
	
	# =====================================================================================
	#                    POS TAG DOCUMENTS
	# =====================================================================================
	spacy_de = spacy.load(args.spacy_model, 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(args.output_file, "w")
	lemmatizer = GermaLemma()
	if args.text_file: write_plain = open(args.text_file, "w")
	
	if ".gz" == args.input_file[-3:]:
		in_file = fu.expand_file(args.input_file)
	else:
		in_file = args.input_file
	
	start = time.time()
	total_processed_sents = 0
	line_generator = fu.file_generator(in_file)
	while file_has_next:
		annos, file_has_next = fu.get_file_annos_chunk(line_generator, chunk_size=CHUNK_SIZE, 		token_class=get_token_type(args.gld_token_type), comment_str=args.comment_str)
		if len(annos) == 0: break
		total_processed_sents += len(annos)
		logger.info(f"Already processed {total_processed_sents} sentences...")
		sents = [a.get_sentence() for a in annos]
		for ix, doc in enumerate(spacy_de.pipe(sents, batch_size=SPACY_BATCH, n_process=SPACY_PROC)):
			conll_str = get_conll_str(annos[ix], doc, use_germalemma=args.use_germalemma)
			write_out.write(conll_str)
			write_out.write("\n\n")
			if args.text_file:
				write_plain.write(" ".join([x.text for x in doc])+"\n")
			
	end = time.time()
	logger.info(f"Processing {args.corpus_name} took {(end - start)} seconds!")