#!/usr/bin/env python3
# coding=utf-8


import argparse
import os
import spacy
import textwrap


parser = argparse.ArgumentParser(description="Tag a .txt file with spaCy and \
                                              save it in .vrt format")
parser.add_argument("-i",
                    dest="input",
                    help="Input file.",
                    required=True)
parser.add_argument("-l",
                    choices=["de", "en", "es", "fr", "pt"],
                    dest="lang",
                    help="Language for tagging",
                    required=True)
parser.add_argument("-o",
                    dest="output",
                    help="Output file.",
                    required=True)
args = parser.parse_args()


SPACY_MODELS = {"de": "de_core_news_sm", "en": "en_core_web_sm",
                "es": "es_core_news_sm", "fr": "fr_core_news_sm",
                "pt": "pt_core_news_sm"}

# Set the language model for spacy
nlp = spacy.load(SPACY_MODELS[args.lang])

# Read text from the input file and if neccessary split it into parts with a
# length of less than 1 million characters.
with open(args.input) as input_file:
    text = input_file.read()
    texts = textwrap.wrap(text, 1000000, break_long_words=False)
    text = None

# Create and open the output file
output_file = open(args.output, "w+")
output_file.write('<?xml version="1.0" encoding="UTF-8"?>\n<corpus>\n<text id="' + os.path.basename(args.input).rsplit(".", 1)[0] + '">\n')
for text in texts:
    # Run spacy nlp over the text (partial string if above 1 million chars)
    doc = nlp(text)
    for sent in doc.sents:
        output_file.write('<s>\n')
        for token in sent:
            # Skip whitespace tokens like "\n" or "\t"
            if token.text.isspace():
                continue
            # Write all information in .vrt style to the output file
            # text, lemma, simple_pos, pos, ner
            output_file.write(token.text + "\t" + token.lemma_ + "\t"
                              + token.pos_ + "\t" + token.tag_ + "\t"
                              + (token.ent_type_ if token.ent_type_ != "" else "NULL") + "\n")
        output_file.write('</s>\n')
output_file.write('</text>\n</corpus>')
output_file.close()