import argparse
import time

def get_arg_parser() -> argparse.ArgumentParser:
    '''
    Create arg parse for active learning training containing options for
        optimizer
        hyper parameters
        model saving
        log directories
    '''
    parser = argparse.ArgumentParser(description='Train  Active Learned Named Entity Recognition on Train Conll2003.')

    # Parser data loader options
    parser.add_argument('--load', action='store_true', help='load the saved files')
    parser.add_argument('--save', action='store_true', help='save the loaded data loaders')
    parser.add_argument('--clean', action='store_true', help='delete saved files')
    parser.add_argument('--train', action='store_true', help='train the network')
    parser.add_argument('--train_elmo_bi_lstm', action='store_true', help='train the ELMo BiLSTM CRF model')
    parser.add_argument('--train_bi_lstm', action='store_true', help='train the BiLSTM CRF model')
    parser.add_argument('--train_dictionary', action='store_true', help='train the dictionary classifier')
    parser.add_argument('--train_phrase_dictionary', action='store_true', help='train the phrase dictionary classifier')
    parser.add_argument('--train_cached', action='store_true', help='train the cached ELMo BiLSTM CRF classifier')
    parser.add_argument('--num_epochs', type=int, default=15, help='number of epochs to train on')
    parser.add_argument('--num_workers', type=int, default=2, help='number of workers for dataset')
    parser.add_argument('--shuffle', type=bool, default=True, help='should shuffle the dataset for training')
    parser.add_argument('--batch_size', type=int, default=1, help='the batch size for specified training')
    parser.add_argument('--embedding_dim', type=int, default=300, help='embedding dimensions for the LSTM features')
    parser.add_argument('--hidden_dim', type=int, default=1000, help='hidden dimensions of LSTM outputs for feature size')
    parser.add_argument('--learning_rate', type=float, default=0.01, help='learning rate for optimizer')
    parser.add_argument('--weight_decay', type=float, default=1e-4, help='weight decay for weight update')
    parser.add_argument('--save_dir', type=str, default='saved_models/', help='directory for saved models')
    parser.add_argument('--model_name', type=str, default='model_{}'.format(time.time()), help='model_name')
    parser.add_argument('--log_dir', type=str, default='tensor_logs/', help='model_name')
    parser.add_argument('--model_path', type=str, help='the path for the saved model')

    parser.add_argument('--sample_size', type=int, default=100, help='the number of samples to use in each active learning iteration')
    parser.add_argument('--oracle_type', type=str, default='simulated', help='the type of oracle to use for the active learning experiment')
    parser.add_argument('--iterations', type=int, default=10, help='The number of active learning iterations to do')
    parser.add_argument('--sampling_strategy', type=str, default='sample', help='The way to select samples either sample or topk')
    parser.add_argument('--dataset', type=str, default='CONLL', help='the supported datasets: SCIERC, CONLL')
    parser.add_argument('--heuristic', type=str, default='uncertain', help='the active leanring heuristic to use')
    parser.add_argument('--cuda', action='store_true', help='use the GPU (CUDA)')
    parser.add_argument('--summary_file', default='summary.csv', type=str, help='name of the summary file being generated by the logger')
    parser.add_argument('--optimizer_type', default='SGD', type=str, help='The optimizer to use')
    parser.add_argument('--momentum', type=float, default=0.0, help='The momentum for SGD optimizer')
    parser.add_argument('--binary_classifier', default=None, type=str, help='if specified, will only classify a certain class')

    return parser