"""
This script allows for training or testing on dev / test of the UD mwt tools.

If run with a single treebank name, it will train or test that treebank.
If run with ud_all or all_ud, it will iterate over all UD treebanks it can find.

Mode can be set to train&dev with --train, to dev set only
with --score_dev, and to test set only with --score_test.

Treebanks are specified as a list.  all_ud or ud_all means to look for
all UD treebanks.

Extra arguments are passed to mwt.  In case the run script
itself is shadowing arguments, you can specify --extra_args as a
parameter to mark where the mwt arguments start.
"""


import io
import logging
import math

from stanza.models import mwt_expander
from stanza.models.common.doc import Document
from stanza.utils.conll import CoNLL
from stanza.utils.training import common
from stanza.utils.training.common import Mode

from stanza.utils.max_mwt_length import max_mwt_length

logger = logging.getLogger('stanza')

def check_mwt(filename):
    """
    Checks whether or not there are MWTs in the given conll file
    """
    doc = CoNLL.conll2doc(filename)
    data = doc.get_mwt_expansions(False)
    return len(data) > 0

def run_treebank(mode, paths, treebank, short_name, command_args, extra_args):
    short_language = short_name.split("_")[0]

    mwt_dir          = paths["MWT_DATA_DIR"]

    train_file       = f"{mwt_dir}/{short_name}.train.in.conllu"
    dev_in_file      = f"{mwt_dir}/{short_name}.dev.in.conllu"
    dev_gold_file    = f"{mwt_dir}/{short_name}.dev.gold.conllu"
    dev_output_file  = f"{mwt_dir}/{short_name}.dev.pred.conllu"
    test_in_file     = f"{mwt_dir}/{short_name}.test.in.conllu"
    test_gold_file   = f"{mwt_dir}/{short_name}.test.gold.conllu"
    test_output_file = f"{mwt_dir}/{short_name}.test.pred.conllu"

    train_json       = f"{mwt_dir}/{short_name}-ud-train-mwt.json"
    dev_json         = f"{mwt_dir}/{short_name}-ud-dev-mwt.json"
    test_json        = f"{mwt_dir}/{short_name}-ud-test-mwt.json"

    eval_file = None
    if '--eval_file' in extra_args:
        eval_file = extra_args[extra_args.index('--eval_file') + 1]

    gold_file = None
    if '--gold_file' in extra_args:
        gold_file = extra_args[extra_args.index('--gold_file') + 1]

    if not check_mwt(train_file):
        logger.info("No training MWTS found for %s.  Skipping" % treebank)
        return
    
    if not check_mwt(dev_in_file) and mode == Mode.TRAIN:
        logger.info("No dev MWTS found for %s.  Training only the deterministic MWT expander" % treebank)
        extra_args.append('--dict_only')

    if mode == Mode.TRAIN:
        max_mwt_len = math.ceil(max_mwt_length([train_json, dev_json]) * 1.1 + 1)
        logger.info("Max len: %f" % max_mwt_len)
        train_args = ['--train_file', train_file,
                      '--eval_file', eval_file if eval_file else dev_in_file,
                      '--gold_file', gold_file if gold_file else dev_gold_file,
                      '--lang', short_language,
                      '--shorthand', short_name,
                      '--mode', 'train',
                      '--max_dec_len', str(max_mwt_len)]
        train_args = train_args + extra_args
        logger.info("Running train step with args: {}".format(train_args))
        mwt_expander.main(train_args)

    if mode == Mode.SCORE_DEV or mode == Mode.TRAIN:
        dev_args = ['--eval_file', eval_file if eval_file else dev_in_file,
                    '--gold_file', gold_file if gold_file else dev_gold_file,
                    '--lang', short_language,
                    '--shorthand', short_name,
                    '--mode', 'predict']
        if command_args.save_output:
            dev_args.extend(['--output_file', dev_output_file])
        dev_args = dev_args + extra_args
        logger.info("Running dev step with args: {}".format(dev_args))
        _, dev_doc = mwt_expander.main(dev_args)
        if not command_args.save_output:
            dev_output_file = "{:C}\n\n".format(dev_doc)
            dev_output_file = io.StringIO(dev_output_file)

        results = common.run_eval_script_mwt(gold_file if gold_file else dev_gold_file, dev_output_file)
        logger.info("Finished running dev set on\n{}\n{}".format(treebank, results))

    if mode == Mode.SCORE_TEST:
        test_args = ['--eval_file', eval_file if eval_file else test_in_file,
                     '--gold_file', gold_file if gold_file else test_gold_file,
                     '--lang', short_language,
                     '--shorthand', short_name,
                     '--mode', 'predict']
        if command_args.save_output:
            test_args.extend(['--output_file', test_output_file])
        test_args = test_args + extra_args
        logger.info("Running test step with args: {}".format(test_args))
        _, test_doc = mwt_expander.main(test_args)
        if not command_args.save_output:
            test_output_file = "{:C}\n\n".format(test_doc)
            test_output_file = io.StringIO(test_output_file)

        results = common.run_eval_script_mwt(gold_file if gold_file else test_gold_file, test_output_file)
        logger.info("Finished running test set on\n{}\n{}".format(treebank, results))

def main():
    common.main(run_treebank, "mwt", "mwt_expander", sub_argparse=mwt_expander.build_argparse())

if __name__ == "__main__":
    main()

