process

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2025 Pavel Vondřička <pavel.vondricka@ff.cuni.cz>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; version 2
# dated June, 1991.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

import os
import sys
from pathlib import Path
import argparse
import configparser
from threading import Thread, Lock
from queue import Queue

from importlib.util import spec_from_loader, module_from_spec
from importlib.machinery import SourceFileLoader 

scriptpath = path = os.path.dirname(os.path.realpath(__file__))+'/'

DEFAULT_THREADS = 10

def load_module(name):
    """
    Load externals scripts as modules
    (they don't the extension `.py`, so they can't be simply imported)
    """
    spec = spec_from_loader(name, SourceFileLoader(name, scriptpath+name))
    module = module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

def progress_bar(iteration, total, prefix='', suffix='', length=30, fill='█'):
    percent = ("{0:.1f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + ' ' * (length - filled_length)
    sys.stderr.write(f'\r{prefix} |{bar}| {percent}% {suffix}\r')
    sys.stderr.flush()

def process_single(config, filename):
    """
    Process single file
    """

    def clear():
        if not config.get('keep_temp_files'):
            txtfile.unlink(missing_ok=True)
            jsonfile.unlink(missing_ok=True)
            annfile.unlink(missing_ok=True)
            annjsonfile.unlink(missing_ok=True)

    # load modules
    xml2standoff = load_module('xml2standoff')
    tagger_ext = config.get('tagger_cmd', None)
    tagger = None if tagger_ext else load_module('tag_'+config.get('tagger', 'ud'))
    ann2standoff = load_module('ann2standoff')
    standoff2xml = load_module('standoff2xml')
    xml2vrt = load_module('xml2vrt')
    
    infile = Path(filename)
    tmp_path = Path(config.get('temp_dir', '/tmp'), infile.name)
    out_path = Path(config.get('output_path', infile.parent), infile.name)
    
    txtfile = tmp_path.with_suffix('.txt')
    jsonfile = tmp_path.with_suffix('.json')
    annfile = tmp_path.with_suffix('.ann')
    annjsonfile = tmp_path.with_suffix('.ann.json')
    annxmlfile = out_path.with_suffix('.ann.xml')

    # load replacements and matchlist for matching analyzer output with original plain text
    replacements = ann2standoff.load_replacements(config)
    matchlist = ann2standoff.load_matchlist(config)

    # process single file
    try:
        xml2standoff.process(config, infile, txtout=txtfile, jsonout=jsonfile)
        if tagger_ext:
            cmd = tagger_ext.replace('{INFILE}', str(txtfile.absolute())).replace('{OUTFILE}', str(annfile.absolute()))
            os.system(cmd)
        else:
            tagger.process(config, txtfile, annfile)
        ann2standoff.process(config, annfile, txtin=txtfile, jsonout=annjsonfile,
                            replacements=replacements, matchlist=matchlist)
        standoff2xml.process(config, txtfile, jsonin=jsonfile, annjsonin=annjsonfile, xmlout=annxmlfile)
        
        if config.get('vrt_path'):
            vrtfile = Path(config.get('vrt_path'), infile.name).with_suffix('.vrt')
            xml2vrt.process(config, annxmlfile, vrtfile)
    except Exception as e:
        sys.stderr.write(infile.name + ': ' + str(e))
        sys.stderr.write("\n")
        sys.stderr.flush()
        if not config.getboolean('keep_failed_files'):
            clear()
        return(False)
    else:
        clear()
        return(True)

counter = 0

def process_batch(config, files, batch_size):
    """
    Process batch of files
    """
    global counter
    counter = 0

    def increment():
        global counter
        counter += 1
        progress_bar(counter, len(files), "Progress:", length=50)

    def processing_worker(q, lock):
        while True:
            file = q.get()
            if file is None:
                break
            succ = process_single(config, file)
            if succ and config.getboolean('progress'):
                lock.acquire()
                increment()
                lock.release()

    queue = Queue()
    lock = Lock()
    # populate the queue
    for file in files:
        queue.put(file)
    workers = [Thread(target=processing_worker, args=(queue, lock)) for _ in range(batch_size)]
    # add a None signal for each worker
    for worker in workers:
        queue.put(None)
    # start all workers
    for worker in workers:
        worker.start()
    # wait for all workers to finish
    for worker in workers:
        worker.join()
    if config.getboolean('progress'):
        sys.stderr.write("\n")
        if counter == len(files):
            sys.stderr.write("All files successfully processed.\n")
        else:
            sys.stderr.write(f"Successfully processed {counter} out of {len(files)} files.\n")
        sys.stderr.flush()


if __name__ == '__main__':
    """
    Pipeline processing
    """
    parser = argparse.ArgumentParser(description="Process XML files by the toolchain")
    parser.add_argument("infiles", help="input XML file names (UTF-8)", nargs='*')
    parser.add_argument("-c", "--config", help="additional config file", type=str)
    parser.add_argument("-p", "--profile", help="config profile to use", type=str, default='DEFAULT')
    parser.add_argument("-m", "--model", help="UD language model to use", type=str)
    parser.add_argument("-t", "--tagger", help="tagger to use (default='ud')", type=str)
    parser.add_argument("-tc", "--tagger-cmd", help="external tagger shell command (overrides -t/--tagger)", type=str)
    parser.add_argument("-O", "--output-path", help="output annotated files to the specified directory", type=str)
    parser.add_argument("-V", "--vrt-path", help="run also xml2vrt and output vertical files to the specified directory", type=str)
    parser.add_argument("-T", "--temp-dir", help="temporary directory to use", type=str)
    parser.add_argument("-W", "--workers", help="number of workers (threads) to use (default=10)", type=int)
    parser.add_argument("-K", "--keep-temp-files", help="keep intermediate (temporary) files (default: delete them)", action="store_true")
    parser.add_argument("-KF", "--keep-failed-files", help="keep intermediate (temporary) files just in case of processing failure", action="store_true")
    parser.add_argument("-P", "--progress", help="show progress bar", action="store_true")
    
    args = parser.parse_args()

    curpath = os.getcwd()
    profiles = configparser.ConfigParser()
    profiles.read([scriptpath+'/xmlanntools.ini', curpath+'/xmlanntools.ini'])
    if args.config:
        read = profiles.read(args.config)
        if read != [args.config]:
            raise Exception("Failed reading configuration file: '{0}'".format(args.config))

    # if no profile was specified, check whether the DEFAULT profile specifies a default profile
    cur_profile = args.profile
    if cur_profile == 'DEFAULT' and profiles[cur_profile].get('profile'):
        cur_profile = profiles[cur_profile].get('profile')

    # evt. update/override currently selected profile configuration with values from command-line arguments
    profiles.read_dict({cur_profile: {k: v for k, v in vars(args).items() if v is not None and v != False}})
    config = profiles[cur_profile]
    
    process_batch(config, args.infiles, config.get('workers', 10))