#!/usr/bin/env python3

import sys
import logging
import argparse


def get_argument_parser():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='action', help='A vanilla recurrent '
                                       'neural network (RNN) language model')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='No logging output except for warnings and errors.')
    # train
    train = subparsers.add_parser('train', help='Train a language model.')
    train.add_argument('data', help='the training data, a plain text file.')
    train.add_argument('-e', '--epochs', type=int, default=10,
                       help='the number of training epochs.')
    train.add_argument('-b', '--batch_size', type=int, default=20)
    train.add_argument('-v', '--vocab_max_size', type=int, default=10000,
                       help='Only keep n most frequent words; treat others as <unk>.')
    train.add_argument('-m', '--save_to', default='model',
                       help='the folder to store model files.')
    train.add_argument('-l', '--log_to', default='logs',
                       help='the folder to store log files. Use for monitoring with tensorboard.')
    # score
    score = subparsers.add_parser('score', help='Use a trained model to score a text.')
    score.add_argument('data', help='the data to be scored, a plain text file.')
    score.add_argument('-m', '--load_from', default='model',
                       help='the folder to load model and vocabulary files from.')
    score.add_argument('-b', '--batch_size', type=int, default=1)
    # sample
    sample = subparsers.add_parser('sample', help='Use a trained model to generate a new text.')
    sample.add_argument('length', type=int, default=100, help='the number of '
                        'symbols to be sampled, i.e., the length of the resulting text.')
    sample.add_argument('-m', '--load_from', default='model',
                        help='the folder to load model and vocabulary files from.')
    sample.add_argument('-f', '--first_symbols', type=lambda s: s.split(" "), default=[],
                        help='the first symbols of the text to be generated.')

    return parser


def action_train(args):
    from romanesco.train import train
    train(**vars(args))


def action_score(args):
    from romanesco.score import score
    perplexity = score(**vars(args))
    print('Perplexity: {0:.2f}'.format(perplexity))


def action_sample(args):
    from romanesco.sample import sample
    generated_text = sample(**vars(args))
    print(generated_text)


def action_generate(args):
    pass


def main():
    parser = get_argument_parser()
    if len(sys.argv) <= 1:
        parser.print_help()
    args = parser.parse_args()
    # log level
    log_level = logging.WARNING if args.quiet else logging.DEBUG
    logging.basicConfig(level=log_level, format='%(asctime)s - %(levelname)s - %(message)s')
    # delegate
    if args.action == 'train':
        action_train(args)
    elif args.action == 'score':
        action_score(args)
    elif args.action == 'sample':
        action_sample(args)


if __name__ == '__main__':
    main()
