# Copyright (C) 2020 Secure Systems Group, University of Waterloo and Aalto University
# License: see LICENSE.txt
# Author: <NAME> and <NAME>

import os
import pandas as pd
import numpy as np
from collections import deque
import tqdm
import joblib

# contains subroutines for training, saving/loading pipeline
import logistic_regression
from sklearn.metrics import classification_report, precision_recall_fscore_support

import data_helpers
from augmentation_helpers import augmentation_helper

# include standard modules
import argparse

# Phase 1: determine experiment parameters
class ExperimentParameters:
    dataset_name = 'threat'
    percentage = 5
    classifier_name = 'char-lr'
    aug_name = 'baseline'
    #aug_name = 'copy'
    augment_classes = [0,1]
    random_seed = 20200303

# lazy imports
def import_optional_libraries():
    if ExperimentParameters.aug_name == 'eda':
        global eda
        from eda_scripts.eda import eda
    elif ExperimentParameters.aug_name == 'add':
        global add_sentences, augmentation_helper_add
        from augmentation_helpers import add_sentences, augmentation_helper_add
    if ExperimentParameters.aug_name == 'ppdb':
        global augment_ppdb
        from ppdb_scripts.augmentation_ppdb import augment_ppdb
    elif ExperimentParameters.aug_name == 'wn':
        global augment_wn
        from wn_scripts.augmentation_wn import augment_wn
    elif ExperimentParameters.aug_name == 'gensim':
        global gensim_aug
        from embedding_scripts.augmentation_gensim import gensim_aug
    elif ExperimentParameters.aug_name == 'bpemb':
        global bpemb_aug
        from embedding_scripts.augmentation_bpemb import bpemb_aug
    elif ExperimentParameters.aug_name == 'gpt2':
        global lm_aug, LanguageModelWrapper, load_lm_corpus
        from gpt2_scripts.augmentation_lm import lm_aug, LanguageModelWrapper, load_lm_corpus

    # import classifier-specific libraries
    if ExperimentParameters.classifier_name == 'bert':
        global BertWrapper
        from fast_bert_scripts.bert_helpers import BertWrapper
    elif ExperimentParameters.classifier_name == 'cnn':
        global CNN_wrapper
        from cnn_scripts.cnn_helpers import CNN_wrapper

def run_experiment():

    import_optional_libraries()

    # Phase 1: determine experiment parameters
    assert ExperimentParameters.dataset_name in ['threat', 'identity_hate']

    ExperimentParameters.aug_dir = '../data/augmentations/%s-%03d/%s-%s'%(ExperimentParameters.dataset_name,
                                                                          ExperimentParameters.percentage,
                                                                          ExperimentParameters.aug_name,
                                                                          ''.join(list(map(str, ExperimentParameters.augment_classes))))

    ExperimentParameters.save_dir = '../results/%s-%03d/%s/%s-%s'%(ExperimentParameters.dataset_name,
                                                               ExperimentParameters.percentage,
                                                               ExperimentParameters.classifier_name,
                                                               ExperimentParameters.aug_name,
                                                               ''.join(list(map(str, ExperimentParameters.augment_classes))))

    ExperimentParameters.results_file = os.path.join(ExperimentParameters.save_dir,'%s.txt'%ExperimentParameters.random_seed)

    if ExperimentParameters.classifier_name == 'bert':
        bert_wrapper = BertWrapper(model_dir = ExperimentParameters.save_dir,
                                  data_dir = ExperimentParameters.aug_dir,
                                  results_filename = '%s'%ExperimentParameters.random_seed,
                                  random_state = ExperimentParameters.random_seed)
    else:
        bert_wrapper = None

    if ExperimentParameters.classifier_name == 'cnn':
        cnn_wrapper = CNN_wrapper(model_dir = ExperimentParameters.save_dir,
                                  data_dir = ExperimentParameters.aug_dir,
                                  results_filename = '%s'%ExperimentParameters.random_seed,
                                  random_state = ExperimentParameters.random_seed)
    else:
        cnn_wrapper = None

    # Phase 2: load datasets
    DATA_DIR = '../data/'

    if not os.path.exists(DATA_DIR):
        os.makedirs(DATA_DIR)

    if not os.path.exists(os.path.join(DATA_DIR,'jigsaw-toxic-comment-classification-challenge')):
        print('Jigsaw dataset not found. ')
        print('Please download from https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge')
        print('.. and extract to the directory data-augmentation/data/jigsaw-toxic-comment-classification-challenge')
        return

    else:
        print('Found Jigsaw dataset.')


    # load train data
    data_dir = '../data/jigsaw-toxic-comment-classification-challenge'
    file = 'train.csv'
    train_csv = os.path.join(data_dir,file)

    assert os.path.exists(train_csv)

    df_train = pd.read_csv(train_csv)


    # load test data
    file = 'test.csv'
    test_csv_comments = os.path.join(data_dir,file)
    df_test_comments = pd.read_csv(test_csv_comments)

    file = 'test_labels.csv'
    test_csv_labels = os.path.join(data_dir,file)
    df_test_labels = pd.read_csv(test_csv_labels)

    # verify 'id' is unique for each row
    assert len(np.unique(df_test_labels['id'].values)) == len(df_test_labels)

    # reindex with 'id'
    df_test_labels.index = df_test_labels.id
    del df_test_labels['id']

    # verify 'id' is unique for each row
    assert len(np.unique(df_test_comments['id'].values)) == len(df_test_comments)

    # reindex with 'id'
    df_test_comments.index = df_test_comments.id
    del df_test_comments['id']

    for ind in df_test_comments.index:
        assert ind in df_test_labels.index

    for ind in df_test_labels.index:
        assert ind in df_test_comments.index

    # merge dataframes (add columns)
    df_test = pd.concat([df_test_comments, df_test_labels], axis=1)
    df_test.head()


    # drop the comments that do not have annotation
    def parse_dataset(df):
        q = '%s != -1'%ExperimentParameters.dataset_name
        df = df.query(q)
        X = df['comment_text'].\
            str.replace('\n',' ').\
            str.replace('\t',' ').\
            str.strip().\
            str.split().\
            str.join(' ')
        y = df[ExperimentParameters.dataset_name]
        return X, y


    X_train, y_train = parse_dataset(df_train)
    X_test, y_test = parse_dataset(df_test)



    # obtain bootstrapped variation of training data
    dataset_name = "%s_%s_%d" % (ExperimentParameters.dataset_name,
                                 'Percentage-%03d'%ExperimentParameters.percentage,
                                 ExperimentParameters.random_seed)


    X_train_tiny, y_train_tiny = data_helpers.get_stratified_subset(X_train, y_train, percentage=ExperimentParameters.percentage, random_state = ExperimentParameters.random_seed)
    data_helpers.save_categorical_dataset(X_train_tiny, y_train_tiny, DATA_DIR, dataset_name)


    if ExperimentParameters.aug_name == 'gpt2':
        # keep all epochs in the same folder
        load_corpus_fun = lambda : load_lm_corpus('../data/'+dataset_name+'.txt')
        print('Creating LMWrapper for dataset')
        model_save_dir = '../models/%s/%s/'%('gpt2', dataset_name)
        lm = LanguageModelWrapper(load_corpus_fun, model_save_dir = model_save_dir, random_state = ExperimentParameters.random_seed)
        if not os.path.exists(model_save_dir):
            print('Fine-tuned GPT2 path not found. Training...')
            lm.train(num_epochs = 2)
        else:
            print('Found fine-tuned GPT2 directory')

    # maps augmentation names to functions
    class Mapper:
        dict = {
            'baseline' : lambda document,num_aug: [document],
            'copy' : lambda document,num_aug: [document]*num_aug,

            # set recommended settings for small-dataset EDA
            'eda' : lambda document,num_aug: eda(document, random_state = ExperimentParameters.random_seed, alpha_sr=0.05, alpha_ri=0.05, alpha_rs=0.05, p_rd=0.05, num_aug=num_aug),

            'add': None,
            'ppdb' : lambda document,num_aug: augment_ppdb(document,num_aug, random_state = ExperimentParameters.random_seed),
            'wn' : lambda document,num_aug: augment_wn(document,num_aug, random_state = ExperimentParameters.random_seed),
            'bpemb' : lambda document,num_aug: bpemb_aug(document,num_aug, random_state = ExperimentParameters.random_seed, max_candidates = 10, incl_orig = True, min_similarity = 0., change_rate = .25),
            'gensim' : lambda document,num_aug: gensim_aug(document,num_aug, random_state = ExperimentParameters.random_seed, max_candidates = 10, incl_orig = True, min_similarity = 0., change_rate = .25),

            'gpt2' : lambda document,num_aug: lm_aug(document, num_aug, random_state = ExperimentParameters.random_seed, incl_orig=True, epoch = 2, lm_wrapper = lm),

            'add_bpemb' : None, # run combine_augmentations.py
            'add_bpemb_gpt2' : None, # run combine_augmentations.py

            # functions for training networks
            # (numpy array, numpy array) -> object
            'char-lr' : lambda x,y: logistic_regression.train_model(x, y, ngram_type='char', random_state = ExperimentParameters.random_seed),
            'word-lr' : lambda x,y: logistic_regression.train_model(x, y, ngram_type='word', random_state = ExperimentParameters.random_seed),
            'bert' : lambda x,y: bert_wrapper.train_model(x, y, num_epochs = 6),
            'cnn': lambda x, y: cnn_wrapper.train_model(x, y, num_epochs = 4)
        }

    assert ExperimentParameters.aug_name in Mapper.dict
    ExperimentParameters.aug_fun = Mapper.dict[ExperimentParameters.aug_name]

    os.makedirs(ExperimentParameters.save_dir, exist_ok=True)

    if os.path.exists(ExperimentParameters.results_file):
        print('Experiment results found.')
        with open(ExperimentParameters.results_file,'r',encoding='utf-8') as f:
            print(f.read())
        print('Exiting...')
        return
    else:
        print('Experiment results will be saved in\n\t',ExperimentParameters.results_file)


    os.makedirs(ExperimentParameters.aug_dir, exist_ok=True)
    ExperimentParameters.aug_file = os.path.join(ExperimentParameters.aug_dir,'%s.txt'%ExperimentParameters.random_seed)

    print('Augmentations will be saved in\n\t',ExperimentParameters.aug_file)

    assert ExperimentParameters.classifier_name in Mapper.dict

    ExperimentParameters.classifier_train_fun = Mapper.dict[ExperimentParameters.classifier_name]
    ExperimentParameters.classifier_file = os.path.join(ExperimentParameters.save_dir,'clf_%s.pkl'%ExperimentParameters.random_seed)

    print('Classifier will be saved in\n\t',ExperimentParameters.classifier_file)


    # run data augmentation
    if not os.path.exists(ExperimentParameters.aug_file):

        if ExperimentParameters.aug_name == 'add':
            augmentation_helper_add(input='../data/%s.txt'%dataset_name, output=ExperimentParameters.aug_file, num_aug = 20, target_classes = ExperimentParameters.augment_classes, random_state = ExperimentParameters.random_seed)
        else:
            augmentation_helper(input='../data/%s.txt'%dataset_name, output=ExperimentParameters.aug_file, aug_function = ExperimentParameters.aug_fun, num_aug = 20, target_classes = ExperimentParameters.augment_classes)


    # load augmented dataset
    X_train_aug, y_train_aug = data_helpers.load_categorical_dataset(ExperimentParameters.aug_file)
    X_train_aug = pd.Series([x.lower() for x in X_train_aug])
    print('\tDataset shapes',X_train_aug.shape, y_train_aug.shape)

    # create classifier
    if not os.path.exists(ExperimentParameters.classifier_file):
        print("Training model...")

        clf = ExperimentParameters.classifier_train_fun(X_train_aug, y_train_aug)

        print("\tSaving model to... %s"%ExperimentParameters.classifier_file)
        joblib.dump(clf, ExperimentParameters.classifier_file)


    print("Reloading Model...")
    clf = joblib.load(ExperimentParameters.classifier_file)

    print(clf.predict_proba(['if you do not stop, the wikapidea nijas will come to your house and kill you']))

    print("Predicting test data...")

    if ExperimentParameters.classifier_name == 'bert':
        y_pred = clf.predict(X_test)
    else:
        prediction_dir = '%s/%s'%(ExperimentParameters.save_dir, 'predictions')
        os.makedirs(prediction_dir, exist_ok=True)
        prediction_file = '%s/%d-output.csv'%(prediction_dir,ExperimentParameters.random_seed)

        print('... saving prediction results to %s.'%prediction_dir)
        y_pred_prob = clf.predict_proba(X_test)
        # output from CNN is a list of lists
        y_pred_prob = np.array(y_pred_prob)
        y_pred = np.array(y_pred_prob[:,1] > 0.5, dtype = int)

        with open(prediction_file, 'w',encoding='utf-8') as f:
            for line in y_pred_prob:
                line = ', '.join(map(str,line))
                f.write('%s\n'%line)


    print(classification_report(y_test, y_pred))

    res_0 = precision_recall_fscore_support(y_test, y_pred, pos_label=0, average='binary')
    res_1 = precision_recall_fscore_support(y_test, y_pred, pos_label=1, average='binary')
    res_m = precision_recall_fscore_support(y_test, y_pred, average='macro')


    with open(ExperimentParameters.results_file,'w',encoding='utf-8') as f:
        f.write('\tPrecision\tRecall\tF-score\n')

        f.write('0:\t')
        f.write('\t'.join(map(str,res_0[:-1])))
        f.write('\n')

        f.write('1:\t')
        f.write('\t'.join(map(str,res_1[:-1])))
        f.write('\n')

        f.write('macro:\t')
        f.write('\t'.join(map(str,res_m[:-1])))
        f.write('\n')


    if args.remove_cache:
        # BERT classifier weigths take 421 MB
        if ExperimentParameters.classifier_name == 'bert':
            bert_wrapper.delete_cache()
            print('Done.')
        # word-lr takes 30 MB per file
        else:
            print('Deleting classifier file...', end='')
            os.remove(ExperimentParameters.classifier_file)
            print('Done.')

if __name__ == '__main__':

    # initiate the parser
    parser = argparse.ArgumentParser()

    # add long and short argument
    parser.add_argument("--dataset_name", "-d", help="obscene/threat", required=True)

    parser.add_argument("--percentage", "-p", help="5", type=int, default=5)

    parser.add_argument("--classifier_name", "-c", help="char-lr/char-word", required=True)

    parser.add_argument("--augmentation_name", "-a", help="baseline/copy/bpemb", required=True)

    parser.add_argument("--random_seed", "-r", help="20200303", type=int, required=True)

    parser.add_argument("--augmentation_classes", "-n", nargs='+', help="0 1 / 1", type=int, default = [1])

    parser.add_argument("--remove_cache", action='store_true')

    # read arguments from the command line
    args = parser.parse_args()

    print()
    print(args)

    # Phase 1: determine experiment parameters
    ExperimentParameters.dataset_name = args.dataset_name
    ExperimentParameters.percentage = args.percentage
    ExperimentParameters.classifier_name = args.classifier_name
    ExperimentParameters.aug_name = args.augmentation_name
    ExperimentParameters.augment_classes = args.augmentation_classes
    ExperimentParameters.random_seed = args.random_seed

    run_experiment()
    print()

# Copyright (C) 2020 Secure Systems Group, University of Waterloo and Aalto University
# License: see LICENSE.txt
# Author: <NAME> and <NAME>

import os
import pandas as pd
import numpy as np
from collections import deque
import tqdm
import joblib

# contains subroutines for training, saving/loading pipeline
import logistic_regression
from sklearn.metrics import classification_report, precision_recall_fscore_support

import data_helpers
from augmentation_helpers import augmentation_helper

# include standard modules
import argparse

# Phase 1: determine experiment parameters
class ExperimentParameters:
    dataset_name = 'threat'
    percentage = 5
    classifier_name = 'char-lr'
    aug_name = 'baseline'
    #aug_name = 'copy'
    augment_classes = [0,1]
    random_seed = 20200303

# lazy imports
def import_optional_libraries():
    if ExperimentParameters.aug_name == 'eda':
        global eda
        from eda_scripts.eda import eda
    elif ExperimentParameters.aug_name == 'add':
        global add_sentences, augmentation_helper_add
        from augmentation_helpers import add_sentences, augmentation_helper_add
    if ExperimentParameters.aug_name == 'ppdb':
        global augment_ppdb
        from ppdb_scripts.augmentation_ppdb import augment_ppdb
    elif ExperimentParameters.aug_name == 'wn':
        global augment_wn
        from wn_scripts.augmentation_wn import augment_wn
    elif ExperimentParameters.aug_name == 'gensim':
        global gensim_aug
        from embedding_scripts.augmentation_gensim import gensim_aug
    elif ExperimentParameters.aug_name == 'bpemb':
        global bpemb_aug
        from embedding_scripts.augmentation_bpemb import bpemb_aug
    elif ExperimentParameters.aug_name == 'gpt2':
        global lm_aug, LanguageModelWrapper, load_lm_corpus
        from gpt2_scripts.augmentation_lm import lm_aug, LanguageModelWrapper, load_lm_corpus

    # import classifier-specific libraries
    if ExperimentParameters.classifier_name == 'bert':
        global BertWrapper
        from fast_bert_scripts.bert_helpers import BertWrapper
    elif ExperimentParameters.classifier_name == 'cnn':
        global CNN_wrapper
        from cnn_scripts.cnn_helpers import CNN_wrapper

def run_experiment():

    import_optional_libraries()

    # Phase 1: determine experiment parameters
    assert ExperimentParameters.dataset_name in ['threat', 'identity_hate']

    ExperimentParameters.aug_dir = '../data/augmentations/%s-%03d/%s-%s'%(ExperimentParameters.dataset_name,
                                                                          ExperimentParameters.percentage,
                                                                          ExperimentParameters.aug_name,
                                                                          ''.join(list(map(str, ExperimentParameters.augment_classes))))

    ExperimentParameters.save_dir = '../results/%s-%03d/%s/%s-%s'%(ExperimentParameters.dataset_name,
                                                               ExperimentParameters.percentage,
                                                               ExperimentParameters.classifier_name,
                                                               ExperimentParameters.aug_name,
                                                               ''.join(list(map(str, ExperimentParameters.augment_classes))))

    ExperimentParameters.results_file = os.path.join(ExperimentParameters.save_dir,'%s.txt'%ExperimentParameters.random_seed)

    if ExperimentParameters.classifier_name == 'bert':
        bert_wrapper = BertWrapper(model_dir = ExperimentParameters.save_dir,
                                  data_dir = ExperimentParameters.aug_dir,
                                  results_filename = '%s'%ExperimentParameters.random_seed,
                                  random_state = ExperimentParameters.random_seed)
    else:
        bert_wrapper = None

    if ExperimentParameters.classifier_name == 'cnn':
        cnn_wrapper = CNN_wrapper(model_dir = ExperimentParameters.save_dir,
                                  data_dir = ExperimentParameters.aug_dir,
                                  results_filename = '%s'%ExperimentParameters.random_seed,
                                  random_state = ExperimentParameters.random_seed)
    else:
        cnn_wrapper = None

    # Phase 2: load datasets
    DATA_DIR = '../data/'

    if not os.path.exists(DATA_DIR):
        os.makedirs(DATA_DIR)

    if not os.path.exists(os.path.join(DATA_DIR,'jigsaw-toxic-comment-classification-challenge')):
        print('Jigsaw dataset not found. ')
        print('Please download from https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge')
        print('.. and extract to the directory data-augmentation/data/jigsaw-toxic-comment-classification-challenge')
        return

    else:
        print('Found Jigsaw dataset.')


    # load train data
    data_dir = '../data/jigsaw-toxic-comment-classification-challenge'
    file = 'train.csv'
    train_csv = os.path.join(data_dir,file)

    assert os.path.exists(train_csv)

    df_train = pd.read_csv(train_csv)


    # load test data
    file = 'test.csv'
    test_csv_comments = os.path.join(data_dir,file)
    df_test_comments = pd.read_csv(test_csv_comments)

    file = 'test_labels.csv'
    test_csv_labels = os.path.join(data_dir,file)
    df_test_labels = pd.read_csv(test_csv_labels)

    # verify 'id' is unique for each row
    assert len(np.unique(df_test_labels['id'].values)) == len(df_test_labels)

    # reindex with 'id'
    df_test_labels.index = df_test_labels.id
    del df_test_labels['id']

    # verify 'id' is unique for each row
    assert len(np.unique(df_test_comments['id'].values)) == len(df_test_comments)

    # reindex with 'id'
    df_test_comments.index = df_test_comments.id
    del df_test_comments['id']

    for ind in df_test_comments.index:
        assert ind in df_test_labels.index

    for ind in df_test_labels.index:
        assert ind in df_test_comments.index

    # merge dataframes (add columns)
    df_test = pd.concat([df_test_comments, df_test_labels], axis=1)
    df_test.head()


    # drop the comments that do not have annotation
    def parse_dataset(df):
        q = '%s != -1'%ExperimentParameters.dataset_name
        df = df.query(q)
        X = df['comment_text'].\
            str.replace('\n',' ').\
            str.replace('\t',' ').\
            str.strip().\
            str.split().\
            str.join(' ')
        y = df[ExperimentParameters.dataset_name]
        return X, y


    X_train, y_train = parse_dataset(df_train)
    X_test, y_test = parse_dataset(df_test)



    # obtain bootstrapped variation of training data
    dataset_name = "%s_%s_%d" % (ExperimentParameters.dataset_name,
                                 'Percentage-%03d'%ExperimentParameters.percentage,
                                 ExperimentParameters.random_seed)


    X_train_tiny, y_train_tiny = data_helpers.get_stratified_subset(X_train, y_train, percentage=ExperimentParameters.percentage, random_state = ExperimentParameters.random_seed)
    data_helpers.save_categorical_dataset(X_train_tiny, y_train_tiny, DATA_DIR, dataset_name)


    if ExperimentParameters.aug_name == 'gpt2':
        # keep all epochs in the same folder
        load_corpus_fun = lambda : load_lm_corpus('../data/'+dataset_name+'.txt')
        print('Creating LMWrapper for dataset')
        model_save_dir = '../models/%s/%s/'%('gpt2', dataset_name)
        lm = LanguageModelWrapper(load_corpus_fun, model_save_dir = model_save_dir, random_state = ExperimentParameters.random_seed)
        if not os.path.exists(model_save_dir):
            print('Fine-tuned GPT2 path not found. Training...')
            lm.train(num_epochs = 2)
        else:
            print('Found fine-tuned GPT2 directory')

    # maps augmentation names to functions
    class Mapper:
        dict = {
            'baseline' : lambda document,num_aug: [document],
            'copy' : lambda document,num_aug: [document]*num_aug,

            # set recommended settings for small-dataset EDA
            'eda' : lambda document,num_aug: eda(document, random_state = ExperimentParameters.random_seed, alpha_sr=0.05, alpha_ri=0.05, alpha_rs=0.05, p_rd=0.05, num_aug=num_aug),

            'add': None,
            'ppdb' : lambda document,num_aug: augment_ppdb(document,num_aug, random_state = ExperimentParameters.random_seed),
            'wn' : lambda document,num_aug: augment_wn(document,num_aug, random_state = ExperimentParameters.random_seed),
            'bpemb' : lambda document,num_aug: bpemb_aug(document,num_aug, random_state = ExperimentParameters.random_seed, max_candidates = 10, incl_orig = True, min_similarity = 0., change_rate = .25),
            'gensim' : lambda document,num_aug: gensim_aug(document,num_aug, random_state = ExperimentParameters.random_seed, max_candidates = 10, incl_orig = True, min_similarity = 0., change_rate = .25),

            'gpt2' : lambda document,num_aug: lm_aug(document, num_aug, random_state = ExperimentParameters.random_seed, incl_orig=True, epoch = 2, lm_wrapper = lm),

            'add_bpemb' : None, # run combine_augmentations.py
            'add_bpemb_gpt2' : None, # run combine_augmentations.py

            # functions for training networks
            # (numpy array, numpy array) -> object
            'char-lr' : lambda x,y: logistic_regression.train_model(x, y, ngram_type='char', random_state = ExperimentParameters.random_seed),
            'word-lr' : lambda x,y: logistic_regression.train_model(x, y, ngram_type='word', random_state = ExperimentParameters.random_seed),
            'bert' : lambda x,y: bert_wrapper.train_model(x, y, num_epochs = 6),
            'cnn': lambda x, y: cnn_wrapper.train_model(x, y, num_epochs = 4)
        }

    assert ExperimentParameters.aug_name in Mapper.dict
    ExperimentParameters.aug_fun = Mapper.dict[ExperimentParameters.aug_name]

    os.makedirs(ExperimentParameters.save_dir, exist_ok=True)

    if os.path.exists(ExperimentParameters.results_file):
        print('Experiment results found.')
        with open(ExperimentParameters.results_file,'r',encoding='utf-8') as f:
            print(f.read())
        print('Exiting...')
        return
    else:
        print('Experiment results will be saved in\n\t',ExperimentParameters.results_file)


    os.makedirs(ExperimentParameters.aug_dir, exist_ok=True)
    ExperimentParameters.aug_file = os.path.join(ExperimentParameters.aug_dir,'%s.txt'%ExperimentParameters.random_seed)

    print('Augmentations will be saved in\n\t',ExperimentParameters.aug_file)

    assert ExperimentParameters.classifier_name in Mapper.dict

    ExperimentParameters.classifier_train_fun = Mapper.dict[ExperimentParameters.classifier_name]
    ExperimentParameters.classifier_file = os.path.join(ExperimentParameters.save_dir,'clf_%s.pkl'%ExperimentParameters.random_seed)

    print('Classifier will be saved in\n\t',ExperimentParameters.classifier_file)


    # run data augmentation
    if not os.path.exists(ExperimentParameters.aug_file):

        if ExperimentParameters.aug_name == 'add':
            augmentation_helper_add(input='../data/%s.txt'%dataset_name, output=ExperimentParameters.aug_file, num_aug = 20, target_classes = ExperimentParameters.augment_classes, random_state = ExperimentParameters.random_seed)
        else:
            augmentation_helper(input='../data/%s.txt'%dataset_name, output=ExperimentParameters.aug_file, aug_function = ExperimentParameters.aug_fun, num_aug = 20, target_classes = ExperimentParameters.augment_classes)


    # load augmented dataset
    X_train_aug, y_train_aug = data_helpers.load_categorical_dataset(ExperimentParameters.aug_file)
    X_train_aug = pd.Series([x.lower() for x in X_train_aug])
    print('\tDataset shapes',X_train_aug.shape, y_train_aug.shape)

    # create classifier
    if not os.path.exists(ExperimentParameters.classifier_file):
        print("Training model...")

        clf = ExperimentParameters.classifier_train_fun(X_train_aug, y_train_aug)

        print("\tSaving model to... %s"%ExperimentParameters.classifier_file)
        joblib.dump(clf, ExperimentParameters.classifier_file)


    print("Reloading Model...")
    clf = joblib.load(ExperimentParameters.classifier_file)

    print(clf.predict_proba(['if you do not stop, the wikapidea nijas will come to your house and kill you']))

    print("Predicting test data...")

    if ExperimentParameters.classifier_name == 'bert':
        y_pred = clf.predict(X_test)
    else:
        prediction_dir = '%s/%s'%(ExperimentParameters.save_dir, 'predictions')
        os.makedirs(prediction_dir, exist_ok=True)
        prediction_file = '%s/%d-output.csv'%(prediction_dir,ExperimentParameters.random_seed)

        print('... saving prediction results to %s.'%prediction_dir)
        y_pred_prob = clf.predict_proba(X_test)
        # output from CNN is a list of lists
        y_pred_prob = np.array(y_pred_prob)
        y_pred = np.array(y_pred_prob[:,1] > 0.5, dtype = int)

        with open(prediction_file, 'w',encoding='utf-8') as f:
            for line in y_pred_prob:
                line = ', '.join(map(str,line))
                f.write('%s\n'%line)


    print(classification_report(y_test, y_pred))

    res_0 = precision_recall_fscore_support(y_test, y_pred, pos_label=0, average='binary')
    res_1 = precision_recall_fscore_support(y_test, y_pred, pos_label=1, average='binary')
    res_m = precision_recall_fscore_support(y_test, y_pred, average='macro')


    with open(ExperimentParameters.results_file,'w',encoding='utf-8') as f:
        f.write('\tPrecision\tRecall\tF-score\n')

        f.write('0:\t')
        f.write('\t'.join(map(str,res_0[:-1])))
        f.write('\n')

        f.write('1:\t')
        f.write('\t'.join(map(str,res_1[:-1])))
        f.write('\n')

        f.write('macro:\t')
        f.write('\t'.join(map(str,res_m[:-1])))
        f.write('\n')


    if args.remove_cache:
        # BERT classifier weigths take 421 MB
        if ExperimentParameters.classifier_name == 'bert':
            bert_wrapper.delete_cache()
            print('Done.')
        # word-lr takes 30 MB per file
        else:
            print('Deleting classifier file...', end='')
            os.remove(ExperimentParameters.classifier_file)
            print('Done.')

if __name__ == '__main__':

    # initiate the parser
    parser = argparse.ArgumentParser()

    # add long and short argument
    parser.add_argument("--dataset_name", "-d", help="obscene/threat", required=True)

    parser.add_argument("--percentage", "-p", help="5", type=int, default=5)

    parser.add_argument("--classifier_name", "-c", help="char-lr/char-word", required=True)

    parser.add_argument("--augmentation_name", "-a", help="baseline/copy/bpemb", required=True)

    parser.add_argument("--random_seed", "-r", help="20200303", type=int, required=True)

    parser.add_argument("--augmentation_classes", "-n", nargs='+', help="0 1 / 1", type=int, default = [1])

    parser.add_argument("--remove_cache", action='store_true')

    # read arguments from the command line
    args = parser.parse_args()

    print()
    print(args)

    # Phase 1: determine experiment parameters
    ExperimentParameters.dataset_name = args.dataset_name
    ExperimentParameters.percentage = args.percentage
    ExperimentParameters.classifier_name = args.classifier_name
    ExperimentParameters.aug_name = args.augmentation_name
    ExperimentParameters.augment_classes = args.augmentation_classes
    ExperimentParameters.random_seed = args.random_seed

    run_experiment()
    print()

# Copyright (C) 2020 Secure Systems Group, University of Waterloo and Aalto University # License: see LICENSE.txt # Author: <NAME> and <NAME> # contains subroutines for training, saving/loading pipeline # include standard modules # Phase 1: determine experiment parameters #aug_name = 'copy' # lazy imports # import classifier-specific libraries # Phase 1: determine experiment parameters # Phase 2: load datasets # load train data # load test data # verify 'id' is unique for each row # reindex with 'id' # verify 'id' is unique for each row # reindex with 'id' # merge dataframes (add columns) # drop the comments that do not have annotation # obtain bootstrapped variation of training data # keep all epochs in the same folder # maps augmentation names to functions # set recommended settings for small-dataset EDA # run combine_augmentations.py # run combine_augmentations.py # functions for training networks # (numpy array, numpy array) -> object # run data augmentation # load augmented dataset # create classifier # output from CNN is a list of lists # BERT classifier weigths take 421 MB # word-lr takes 30 MB per file # initiate the parser # add long and short argument # read arguments from the command line # Phase 1: determine experiment parameters