import argparse
import multiprocessing as mp
import os
import xml.etree.ElementTree as et
from io import StringIO

import nltk

buggy_annotations = [6538624, 6550700, 6547918, 6521702, 6541530, 6774318, 4531088, 4531238]
def iterate_docs(folder_path, tmp_path):
    os.makedirs(tmp_path, exist_ok=True)
    ft_path = os.path.join(folder_path, 'fulltext')
    all_docs = list(filter(lambda x: x.endswith('.xml'), os.listdir(ft_path)))
    for doc_name in all_docs:
        it = et.iterparse(StringIO(open(os.path.join(ft_path, doc_name)).read()))
        for _, el in it:
            prefix, has_namespace, postfix = el.tag.partition('}')
            if has_namespace:
                el.tag = postfix
        root = it.root
        for sentence in root:
            for annotation in sentence:
                if "ID" in annotation.attrib and int(annotation.attrib['ID']) in buggy_annotations:
                    print('Delete one buggy annotation from', doc_name)
                    sentence.remove(annotation)
                    break
        dump_path = os.path.join(tmp_path, doc_name)
        et.ElementTree(root).write(dump_path)
        # doc_xml.write(dump_path, default_namespace='')
        yield dump_path


def process_doc(script_folder, doc_path):
    os.chdir(script_folder)
    print(f'processing {doc_path}...')
    cmd = f'perl fttosem.pl {doc_path}'
    print(cmd)
    os.system(cmd)
    print('Done')
    return True


def main():
    parser = argparse.ArgumentParser('Run fttosem perl examples.')
    parser.add_argument('-s', help='script folder', type=str, required=True)
    parser.add_argument('-p', help='path to corpora', type=str)
    parser.add_argument('-o', help='output path', type=str, default='/tmp/framenet')
    args = parser.parse_args()
    script_folder = args.s
    corpora_folder = args.p
    if corpora_folder is None:
        nltk.download('framenet')
        corpora_folder = os.path.join(nltk.data.path[0], 'corpora', 'framenet_v17')
    os.chdir(script_folder)
    fns = list(iterate_docs(corpora_folder, args.o))
    print(f'{len(fns)} documents detected.')

    processes = list()
    for fn in fns:
        print('Processing', fn)
        process = mp.Process(target=process_doc, args=(script_folder, fn))
        process.start()
        processes.append(process)
    for process in processes:
        process.join(timeout=480)

    rst = os.listdir(args.o)
    rst = list(filter(lambda x: x.endswith('.sem'), rst))
    print(f'{len(rst)} Done.')
    rst = [fn[:-4] for fn in rst]
    fns = [fn[:-4] for fn in fns]
    print('Unfinished docs:')
    for fn in set(fns) - set(rst):
        print(fn)


if __name__ == '__main__':
    main()