import argparse
import glob
import logging
import os
import pathlib
import random
import sys

from enum import Enum

from stanza.resources.default_packages import default_charlms, lemma_charlms, tokenizer_charlms, pos_charlms, depparse_charlms, TRANSFORMERS, TRANSFORMER_LAYERS
from stanza.resources.default_packages import no_pretrain_languages, pos_pretrains, depparse_pretrains, default_pretrains
from stanza.models.common.constant import treebank_to_short_name
from stanza.models.common.utils import ud_scores
from stanza.resources.common import download, DEFAULT_MODEL_DIR, UnknownLanguageError
from stanza.utils.datasets import common
import stanza.utils.default_paths as default_paths
from stanza.utils import conll18_ud_eval as ud_eval

logger = logging.getLogger('stanza')

class Mode(Enum):
    TRAIN = 1
    SCORE_DEV = 2
    SCORE_TEST = 3
    SCORE_TRAIN = 4

class ArgumentParserWithExtraHelp(argparse.ArgumentParser):
    def __init__(self, sub_argparse, *args, **kwargs):
        super().__init__(*args, **kwargs)  # forwards all unused arguments

        self.sub_argparse = sub_argparse

    def print_help(self, file=None):
        super().print_help(file=file)

    def format_help(self):
        help_text = super().format_help()
        if self.sub_argparse is not None:
            sub_text = self.sub_argparse.format_help().split("\n")
            first_line = -1
            for line_idx, line in enumerate(sub_text):
                if line.strip().startswith("usage:"):
                    first_line = line_idx
                elif first_line >= 0 and not line.strip():
                    first_line = line_idx
                    break
            help_text = help_text + "\n\nmodel arguments:" + "\n".join(sub_text[first_line:])
        return help_text


def build_argparse(sub_argparse=None):
    parser = ArgumentParserWithExtraHelp(sub_argparse=sub_argparse, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--save_output', dest='save_output', default=False, action='store_true', help="Save output - default is to use a temp directory.")

    parser.add_argument('treebanks', type=str, nargs='+', help='Which treebanks to run on.  Use all_ud or ud_all for all UD treebanks')

    parser.add_argument('--train', dest='mode', default=Mode.TRAIN, action='store_const', const=Mode.TRAIN, help='Run in train mode')
    parser.add_argument('--score_dev', dest='mode', action='store_const', const=Mode.SCORE_DEV, help='Score the dev set')
    parser.add_argument('--score_test', dest='mode', action='store_const', const=Mode.SCORE_TEST, help='Score the test set')
    parser.add_argument('--score_train', dest='mode', action='store_const', const=Mode.SCORE_TRAIN, help='Score the train set as a test set.  Currently only implemented for some models')

    # These arguments need to be here so we can identify if the model already exists in the user-specified home
    # TODO: when all of the model scripts handle their own names, can eliminate this argument
    parser.add_argument('--save_dir', type=str, default=None, help="Root dir for saving models.  If set, will override the model's default.")
    parser.add_argument('--save_name', type=str, default=None, help="Base name for saving models.  If set, will override the model's default.")

    parser.add_argument('--charlm_only', action='store_true', default=False, help='When asking for ud_all, filter the ones which have charlms')
    parser.add_argument('--transformer_only', action='store_true', default=False, help='When asking for ud_all, filter the ones for languages where we have transformers')

    parser.add_argument('--force', dest='force', action='store_true', default=False, help='Retrain existing models')
    return parser

def add_charlm_args(parser):
    parser.add_argument('--charlm', default="default", type=str, help='Which charlm to run on.  Will use the default charlm for this language/model if not set.  Set to None to turn off charlm for languages with a default charlm')
    parser.add_argument('--no_charlm', dest='charlm', action="store_const", const=None, help="Don't use a charlm, even if one is used by default for this package")

def main(run_treebank, model_dir, model_name, add_specific_args=None, sub_argparse=None, build_model_filename=None, choose_charlm_method=None, args=None):
    """
    A main program for each of the run_xyz scripts

    It collects the arguments and runs the main method for each dataset provided.
    It also tries to look for an existing model and not overwrite it unless --force is provided

    model_name can be a callable expecting the args
      - the charlm, for example, needs this feature, since it makes
        both forward and backward models
    """
    if args is None:
        logger.info("Training program called with:\n" + " ".join(sys.argv))
        args = sys.argv[1:]
    else:
        logger.info("Training program called with:\n" + " ".join(args))

    paths = default_paths.get_default_paths()

    parser = build_argparse(sub_argparse)
    if add_specific_args is not None:
        add_specific_args(parser)
    if '--extra_args' in sys.argv:
        idx = sys.argv.index('--extra_args')
        extra_args = sys.argv[idx+1:]
        command_args = parser.parse_args(sys.argv[:idx])
    else:
        command_args, extra_args = parser.parse_known_args(args=args)

    # Pass this through to the underlying model as well as use it here
    # we don't put --save_name here for the awkward situation of
    # --save_name being specified for an invocation with multiple treebanks
    if command_args.save_dir:
        extra_args.extend(["--save_dir", command_args.save_dir])

    # if --no_seed is added to the args, we actually want to pick a seed here
    # that way, save file names will be consistent...
    # otherwise, it might try to use different save names when using the
    # train and dev sets, if the random seed is used as part of the save name
    while '--no_seed' in extra_args:
        idx = extra_args.index('--no_seed')
        random_seed = random.randint(0, 1000000000)
        logger.info("Using random seed %d", random_seed)
        extra_args[idx:idx+1] = ['--seed', str(random_seed)]

    if callable(model_name):
        model_name = model_name(command_args)

    mode = command_args.mode
    treebanks = []

    for treebank in command_args.treebanks:
        # this is a really annoying typo to make if you copy/paste a
        # UD directory name on the cluster and your job dies 30s after
        # being queued for an hour
        if treebank.endswith("/"):
            treebank = treebank[:-1]
        if treebank.lower() in ('ud_all', 'all_ud'):
            ud_treebanks = common.get_ud_treebanks(paths["UDBASE"])
            if choose_charlm_method is not None and command_args.charlm_only:
                logger.info("Filtering ud_all treebanks to only those which can use charlm for this model")
                ud_treebanks = [x for x in ud_treebanks
                                if choose_charlm_method(*treebank_to_short_name(x).split("_", 1), 'default') is not None]
            if command_args.transformer_only:
                logger.info("Filtering ud_all treebanks to only those which can use a transformer for this model")
                ud_treebanks = [x for x in ud_treebanks if treebank_to_short_name(x).split("_")[0] in TRANSFORMERS]
            logger.info("Expanding %s to %s", treebank, " ".join(ud_treebanks))
            treebanks.extend(ud_treebanks)
        else:
            treebanks.append(treebank)

    for treebank_idx, treebank in enumerate(treebanks):
        if treebank_idx > 0:
            logger.info("=========================================")

        short_name = treebank_to_short_name(treebank)
        logger.debug("%s: %s" % (treebank, short_name))

        save_name_args = []
        if model_name != 'ete':
            # ete is several models at once, so we don't set --save_name
            # theoretically we could handle a parametrized save_name
            if command_args.save_name:
                save_name = command_args.save_name
                # if there's more than 1 treebank, we can't save them all to this save_name
                # we have to override that value for each treebank
                if len(treebanks) > 1:
                    save_name_dir, save_name_filename = os.path.split(save_name)
                    save_name_filename = "%s_%s" % (short_name, save_name_filename)
                    save_name = os.path.join(save_name_dir, save_name_filename)
                    logger.info("Save file for %s model for %s: %s", short_name, treebank, save_name)
                save_name_args = ['--save_name', save_name]
            # some run scripts can build the model filename
            # in order to check for models that are already created
            elif build_model_filename is None:
                save_name = "%s_%s.pt" % (short_name, model_name)
                logger.info("Save file for %s model: %s", short_name, save_name)
                save_name_args = ['--save_name', save_name]
            else:
                save_name_args = []

            if mode == Mode.TRAIN and not command_args.force:
                if build_model_filename is not None:
                    model_path = build_model_filename(paths, short_name, command_args, extra_args)
                elif command_args.save_dir:
                    model_path = os.path.join(command_args.save_dir, save_name)
                else:
                    save_dir = os.path.join("saved_models", model_dir)
                    save_name_args.extend(["--save_dir", save_dir])
                    model_path = os.path.join(save_dir, save_name)

                if model_path is None:
                    # this can happen with the identity lemmatizer, for example
                    pass
                elif os.path.exists(model_path):
                    logger.info("%s: %s exists, skipping!" % (treebank, model_path))
                    continue
                else:
                    logger.info("%s: %s does not exist, training new model" % (treebank, model_path))

        run_treebank(mode, paths, treebank, short_name, command_args, extra_args + save_name_args)

def run_eval_script(gold_conllu_file, system_conllu_file, evals=None):
    """ Wrapper for lemma scorer. """
    evaluation = ud_scores(gold_conllu_file, system_conllu_file)

    if evals is None:
        return ud_eval.build_evaluation_table(evaluation, verbose=True, counts=False, enhanced=False)
    else:
        results = [evaluation[key].f1 for key in evals]
        max_len = max(5, max(len(e) for e in evals))
        evals_string = " ".join(("{:>%d}" % max_len).format(e) for e in evals)
        results_string = " ".join(("{:%d.2f}" % max_len).format(100 * x) for x in results)
        return evals_string + "\n" + results_string

def run_eval_script_tokens(eval_gold, eval_pred):
    return run_eval_script(eval_gold, eval_pred, evals=["Tokens", "Sentences", "Words"])

def run_eval_script_mwt(eval_gold, eval_pred):
    return run_eval_script(eval_gold, eval_pred, evals=["Words"])

def run_eval_script_pos(eval_gold, eval_pred):
    return run_eval_script(eval_gold, eval_pred, evals=["UPOS", "XPOS", "UFeats", "AllTags"])

def run_eval_script_depparse(eval_gold, eval_pred):
    return run_eval_script(eval_gold, eval_pred, evals=["UAS", "LAS", "CLAS", "MLAS", "BLEX"])


def find_wordvec_pretrain(language, default_pretrains, dataset_pretrains=None, dataset=None, model_dir=DEFAULT_MODEL_DIR):
    # try to get the default pretrain for the language,
    # but allow the package specific value to override it if that is set
    default_pt = default_pretrains.get(language, None)
    if dataset is not None and dataset_pretrains is not None:
        default_pt = dataset_pretrains.get(language, {}).get(dataset, default_pt)

    if default_pt is not None:
        default_pt_path = '{}/{}/pretrain/{}.pt'.format(model_dir, language, default_pt)
        if not os.path.exists(default_pt_path):
            logger.info("Default pretrain should be {}  Attempting to download".format(default_pt_path))
            try:
                download(lang=language, package=None, processors={"pretrain": default_pt}, model_dir=model_dir)
            except UnknownLanguageError:
                # if there's a pretrain in the directory, hiding this
                # error will let us find that pretrain later
                pass
        if os.path.exists(default_pt_path):
            if dataset is not None and dataset_pretrains is not None and language in dataset_pretrains and dataset in dataset_pretrains[language]:
                logger.info(f"Using default pretrain for {language}:{dataset}, found in {default_pt_path}  To use a different pretrain, specify --wordvec_pretrain_file")
            else:
                logger.info(f"Using default pretrain for language, found in {default_pt_path}  To use a different pretrain, specify --wordvec_pretrain_file")
            return default_pt_path

    pretrain_path = '{}/{}/pretrain/*.pt'.format(model_dir, language)
    pretrains = glob.glob(pretrain_path)
    if len(pretrains) == 0:
        # we already tried to download the default pretrain once
        # and it didn't work.  maybe the default language package
        # will have something?
        logger.warning(f"Cannot figure out which pretrain to use for '{language}'.  Will download the default package and hope for the best")
        try:
            download(lang=language, model_dir=model_dir)
        except UnknownLanguageError as e:
            # this is a very unusual situation
            # basically, there was a language which we started to add
            # to the resources, but then didn't release the models
            # as part of resources.json
            raise FileNotFoundError(f"Cannot find any pretrains in {pretrain_path}  No pretrains in the system for this language.  Please prepare an embedding as a .pt and use --wordvec_pretrain_file to specify a .pt file to use") from e
        pretrains = glob.glob(pretrain_path)
    if len(pretrains) == 0:
        raise FileNotFoundError(f"Cannot find any pretrains in {pretrain_path}  Try 'stanza.download(\"{language}\")' to get a default pretrain or use --wordvec_pretrain_file to specify a .pt file to use")
    if len(pretrains) > 1:
        raise FileNotFoundError(f"Too many pretrains to choose from in {pretrain_path}  Must specify an exact path to a --wordvec_pretrain_file")
    pt = pretrains[0]
    logger.info(f"Using pretrain found in {pt}  To use a different pretrain, specify --wordvec_pretrain_file")
    return pt

def find_charlm_file(direction, language, charlm, model_dir=DEFAULT_MODEL_DIR):
    """
    Return the path to the forward or backward charlm if it exists for the given package

    If we can figure out the package, but can't find it anywhere, we try to download it
    """
    saved_path = 'saved_models/charlm/{}_{}_{}_charlm.pt'.format(language, charlm, direction)
    if os.path.exists(saved_path):
        logger.info(f'Using model {saved_path} for {direction} charlm')
        return saved_path

    resource_path = '{}/{}/{}_charlm/{}.pt'.format(model_dir, language, direction, charlm)
    if os.path.exists(resource_path):
        logger.info(f'Using model {resource_path} for {direction} charlm')
        return resource_path

    try:
        download(lang=language, package=None, processors={f"{direction}_charlm": charlm}, model_dir=model_dir)
        if os.path.exists(resource_path):
            logger.info(f'Downloaded model, using model {resource_path} for {direction} charlm')
            return resource_path
    except ValueError as e:
        raise FileNotFoundError(f"Cannot find {direction} charlm in either {saved_path} or {resource_path}  Attempted downloading {charlm} but that did not work") from e

    raise FileNotFoundError(f"Cannot find {direction} charlm in either {saved_path} or {resource_path}  Attempted downloading {charlm} but that did not work")

def build_charlm_args(language, charlm, base_args=True, model_dir=DEFAULT_MODEL_DIR, use_backward_model=True):
    """
    If specified, return forward and backward charlm args
    """
    if charlm:
        try:
            forward = find_charlm_file('forward', language, charlm, model_dir=model_dir)
            if use_backward_model:
                backward = find_charlm_file('backward', language, charlm, model_dir=model_dir)
        except FileNotFoundError as e:
            # if we couldn't find sd_isra when training an SD model,
            # for example, but isra exists, we try to download the
            # shorter model name
            if charlm.startswith(language + "_"):
                short_charlm = charlm[len(language)+1:]
                try:
                    forward = find_charlm_file('forward', language, short_charlm, model_dir=model_dir)
                    if use_backward_model:
                        backward = find_charlm_file('backward', language, short_charlm, model_dir=model_dir)
                except FileNotFoundError as e2:
                    raise FileNotFoundError("Tried to find charlm %s, which doesn't exist.  Also tried %s, but didn't find that either" % (charlm, short_charlm)) from e
                logger.warning("Was asked to find charlm %s, which does not exist.  Did find %s though", charlm, short_charlm)
            else:
                raise

        char_args = ['--charlm_forward_file', forward]
        if use_backward_model:
            char_args += ['--charlm_backward_file', backward]
        if not base_args:
            return char_args
        return ['--charlm',
                '--charlm_shorthand', f'{language}_{charlm}'] + char_args

    return []

def choose_charlm(language, dataset, charlm, language_charlms, dataset_charlms):
    """
    charlm == "default" means the default charlm for this dataset or language
    charlm == None is no charlm
    """
    default_charlm = language_charlms.get(language, None)
    specific_charlm = dataset_charlms.get(language, {}).get(dataset, None)

    if charlm is None:
        return None
    elif charlm != "default":
        return charlm
    elif dataset in dataset_charlms.get(language, {}):
        # this way, a "" or None result gets honored
        # thus treating "not in the map" as a way for dataset_charlms to signal to use the default
        return specific_charlm
    elif default_charlm:
        return default_charlm
    else:
        return None

def choose_pos_charlm(short_language, dataset, charlm):
    """
    charlm == "default" means the default charlm for this dataset or language
    charlm == None is no charlm
    """
    return choose_charlm(short_language, dataset, charlm, default_charlms, pos_charlms)

def choose_depparse_charlm(short_language, dataset, charlm):
    """
    charlm == "default" means the default charlm for this dataset or language
    charlm == None is no charlm
    """
    return choose_charlm(short_language, dataset, charlm, default_charlms, depparse_charlms)

def choose_lemma_charlm(short_language, dataset, charlm):
    """
    charlm == "default" means the default charlm for this dataset or language
    charlm == None is no charlm
    """
    return choose_charlm(short_language, dataset, charlm, default_charlms, lemma_charlms)

def choose_tokenizer_charlm(short_language, dataset, charlm):
    """
    charlm == "default" means the default charlm for this dataset or language
    charlm == None is no charlm
    """
    return choose_charlm(short_language, dataset, charlm, default_charlms, tokenizer_charlms)

def choose_transformer(short_language, command_args, extra_args, warn=True, layers=False):
    """
    Choose a transformer using the default options for this language
    """
    bert_args = []
    if command_args is not None and command_args.use_bert and '--bert_model' not in extra_args:
        if short_language in TRANSFORMERS:
            bert_args = ['--bert_model', TRANSFORMERS.get(short_language)]
            if layers and short_language in TRANSFORMER_LAYERS and '--bert_hidden_layers' not in extra_args:
                bert_args.extend(['--bert_hidden_layers', str(TRANSFORMER_LAYERS.get(short_language))])
        elif warn:
            logger.error("Transformer requested, but no default transformer for %s  Specify one using --bert_model" % short_language)

    return bert_args

def build_pos_charlm_args(short_language, dataset, charlm, base_args=True, model_dir=DEFAULT_MODEL_DIR):
    charlm = choose_pos_charlm(short_language, dataset, charlm)
    charlm_args = build_charlm_args(short_language, charlm, base_args, model_dir)
    return charlm_args

def build_lemma_charlm_args(short_language, dataset, charlm, base_args=True, model_dir=DEFAULT_MODEL_DIR):
    charlm = choose_lemma_charlm(short_language, dataset, charlm)
    charlm_args = build_charlm_args(short_language, charlm, base_args, model_dir)
    return charlm_args

def build_depparse_charlm_args(short_language, dataset, charlm, base_args=True, model_dir=DEFAULT_MODEL_DIR):
    charlm = choose_depparse_charlm(short_language, dataset, charlm)
    charlm_args = build_charlm_args(short_language, charlm, base_args, model_dir)
    return charlm_args

def build_tokenizer_charlm_args(short_language, dataset, charlm, base_args=True, model_dir=DEFAULT_MODEL_DIR):
    charlm = choose_tokenizer_charlm(short_language, dataset, charlm)
    charlm_args = build_charlm_args(short_language, charlm, base_args, model_dir, use_backward_model=False)
    return charlm_args


def build_wordvec_args(short_language, dataset, extra_args, task_pretrains):
    if '--wordvec_pretrain_file' in extra_args or '--no_pretrain' in extra_args:
        return []

    if short_language in no_pretrain_languages:
        # we couldn't find word vectors for a few languages...:
        # coptic, naija, old russian, turkish german, swedish sign language
        logger.warning("No known word vectors for language {}  If those vectors can be found, please update the training scripts.".format(short_language))
        return ["--no_pretrain"]
    else:
        if short_language in task_pretrains and dataset in task_pretrains[short_language]:
            dataset_pretrains = task_pretrains
        else:
            dataset_pretrains = {}
        wordvec_pretrain = find_wordvec_pretrain(short_language, default_pretrains, dataset_pretrains, dataset)
        return ["--wordvec_pretrain_file", wordvec_pretrain]

def build_pos_wordvec_args(short_language, dataset, extra_args):
    return build_wordvec_args(short_language, dataset, extra_args, pos_pretrains)

def build_depparse_wordvec_args(short_language, dataset, extra_args):
    return build_wordvec_args(short_language, dataset, extra_args, depparse_pretrains)

