import argparse

from utils import read_data, save_pickle, read_ade_data
from biobert_ner.utils_ner import generate_input_files
from biobert_re.utils_re import generate_re_input_files
from typing import List, Iterator, Dict
import warnings
import os
import re

labels = ['B-DRUG', 'I-DRUG', 'B-STR', 'I-STR', 'B-DUR', 'I-DUR',
          'B-ROU', 'I-ROU', 'B-FOR', 'I-FOR', 'B-ADE', 'I-ADE',
          'B-DOS', 'I-DOS', 'B-REA', 'I-REA', 'B-FRE', 'I-FRE', 'O']


def parse_arguments():
    """Parses program arguments"""
    parser = argparse.ArgumentParser()
    parser.add_argument("--task", type=str,
                        help="Task to be completed. 'NER', 'RE'. Default is 'NER'.",
                        default="NER")

    parser.add_argument("--input_dir", type=str,
                        help="Directory with txt and ann files. Default is 'data/'.",
                        default="data/")

    parser.add_argument("--ade_dir", type=str,
                        help="Directory with ADE corpus. Default is None.",
                        default=None)

    parser.add_argument("--target_dir", type=str,
                        help="Directory to save files. Default is 'dataset/'.",
                        default='dataset/')

    parser.add_argument("--max_seq_len", type=int,
                        help="Maximum sequence length. Default is 512.",
                        default=512)

    parser.add_argument("--dev_split", type=float,
                        help="Ratio of dev data. Default is 0.1",
                        default=0.1)

    parser.add_argument("--tokenizer", type=str,
                        help="The tokenizer to use. 'scispacy', 'scispacy_plus', 'biobert-base', 'biobert-large', 'default'.",
                        default="scispacy")

    parser.add_argument("--ext", type=str,
                        help="Extension of target file. Default is txt.",
                        default="txt")

    parser.add_argument("--sep", type=str,
                        help="Token-label separator. Default is a space.",
                        default=" ")

    arguments = parser.parse_args()
    return arguments


def default_tokenizer(sequence: str) -> List[str]:
    """A tokenizer that splits sequence by a whitespace."""
    words = re.split("\n| |\t", sequence)
    tokens = []
    for word in words:
        word = word.strip()

        if not word:
            continue

        tokens.append(word)

    return tokens


def scispacy_plus_tokenizer(sequence: str, scispacy_tok=None) -> Iterator[str]:
    """
    Runs the scispacy tokenizer and removes all tokens with
    just whitespace characters
    """
    if scispacy_tok is None:
        import en_ner_bc5cdr_md
        scispacy_tok = en_ner_bc5cdr_md.load().tokenizer

    scispacy_tokens = list(map(lambda x: str(x), scispacy_tok(sequence)))
    tokens = filter(lambda t: not (' ' in t or '\n' in t or '\t' in t), scispacy_tokens)

    return tokens


def ner_generator(files: Dict[str, tuple], args) -> None:
    """Generates files for NER"""
    # Generate train, dev, test files
    for filename, data in files.items():
        generate_input_files(ehr_records=data[0], ade_records=data[1],
                             filename=args.target_dir + filename + '.' + args.ext,
                             max_len=args.max_seq_len, sep=args.sep)
        save_pickle(args.target_dir + filename, {"EHR": data[0], "ADE": data[1]})

    # Generate labels file
    with open(args.target_dir + 'labels.txt', 'w') as file:
        output_labels = map(lambda x: x + '\n', labels)
        file.writelines(output_labels)

    filenames = [name for files in map(
        lambda x: [x + '.' + args.ext, x + '.pkl'],
        list(files.keys()))
                 for name in files]

    print("\nGenerating files successful. Files generated: ",
          ', '.join(filenames), ', labels.txt', sep='')


def re_generator(files: Dict[str, tuple], args):
    """Generates files for RE"""
    for filename, data in files.items():
        generate_re_input_files(ehr_records=data[0], ade_records=data[1],
                                filename=args.target_dir + filename + '.' + args.ext,
                                max_len=args.max_seq_len, sep=args.sep,
                                is_test=data[2], is_label=data[3])

    save_pickle(args.target_dir + 'train', {"EHR": files['train'][0], "ADE": files['train'][1]})
    save_pickle(args.target_dir + 'test',  {"EHR": files['test'][0],  "ADE": files['test'][1]})

    print("\nGenerating files successful. Files generated: ",
          'train.tsv,', 'dev.tsv,', 'test.tsv,',
          'test_labels.tsv,', 'train_rel.pkl,', 'test_rel.pkl,', 'test_labels_rel.pkl', sep=' ')


def main():
    args = parse_arguments()

    if args.target_dir[-1] != '/':
        args.target_dir += '/'

    if args.sep == "tab":
        args.sep = "\t"

    if not os.path.isdir(args.target_dir):
        os.mkdir(args.target_dir)

    if args.tokenizer == "default":
        tokenizer = default_tokenizer
        is_bert_tokenizer = False

    elif args.tokenizer == "scispacy":
        import en_ner_bc5cdr_md
        tokenizer = en_ner_bc5cdr_md.load().tokenizer
        is_bert_tokenizer = False

    elif args.tokenizer == 'scispacy_plus':
        import en_ner_bc5cdr_md
        scispacy_tok = en_ner_bc5cdr_md.load().tokenizer
        scispacy_plus_tokenizer.__defaults__ = (scispacy_tok,)

        tokenizer = scispacy_plus_tokenizer
        is_bert_tokenizer = False

    elif args.tokenizer == 'biobert-large':
        from transformers import AutoTokenizer
        biobert = AutoTokenizer.from_pretrained(
            "dmis-lab/biobert-large-cased-v1.1")

        args.max_seq_len -= biobert.num_special_tokens_to_add()
        tokenizer = biobert.tokenize
        is_bert_tokenizer = True

    elif args.tokenizer == 'biobert-base':
        from transformers import AutoTokenizer
        biobert = AutoTokenizer.from_pretrained(
            "dmis-lab/biobert-base-cased-v1.1")

        args.max_seq_len -= biobert.num_special_tokens_to_add()
        tokenizer = biobert.tokenize
        is_bert_tokenizer = True

    else:
        warnings.warn("Tokenizer named " + args.tokenizer + " not found."
                      "Using default tokenizer instead. Acceptable values"
                      "include 'scispacy', 'biobert-base', 'biobert-large',"
                      "and 'default'.")
        tokenizer = default_tokenizer
        is_bert_tokenizer = False

    print("\nReading data\n")
    train_dev, test = read_data(data_dir=args.input_dir,
                                tokenizer=tokenizer,
                                is_bert_tokenizer=is_bert_tokenizer,
                                verbose=1)

    if args.ade_dir is not None:
        ade_train_dev = read_ade_data(ade_data_dir=args.ade_dir, verbose=1)

        ade_dev_split_idx = int((1 - args.dev_split) * len(ade_train_dev))
        ade_train = ade_train_dev[:ade_dev_split_idx]
        ade_devel = ade_train_dev[ade_dev_split_idx:]

    else:
        ade_train_dev = None
        ade_train = None
        ade_devel = None

    print('\n')

    # Data is already shuffled, just split for dev set
    dev_split_idx = int((1 - args.dev_split) * len(train_dev))
    train = train_dev[:dev_split_idx]
    devel = train_dev[dev_split_idx:]

    # Data for NER
    if args.task.lower() == 'ner':
        files = {'train': (train, ade_train), 'train_dev': (train_dev, ade_train_dev),
                 'devel': (devel, ade_devel), 'test': (test, None)}

        ner_generator(files, args)

    # Data for RE
    elif args.task.lower() == 're':
        # {dataset_name: (ehr_data, ade_data, is_test, is_label)}
        files = {'train': (train, ade_train, False, True), 'dev': (devel, ade_devel, False, True),
                 'test': (test, None, True, False), 'test_labels': (test, None, True, True)}

        re_generator(files, args)


if __name__ == '__main__':
    main()