import argparse
from datetime import datetime

class Parser(object):
    def __init__(self):
        parser = argparse.ArgumentParser()

        # Parameters for Hyper-param sweep
        parser.add_argument("--base", default=-1, help="Base counter for Hyper-param search", type=int)
        parser.add_argument("--inc", default=-1, help="Increment counter for Hyper-param search", type=int)
        parser.add_argument("--hyper", default=-1, help="Which Hyper param settings", type=int)
        parser.add_argument("--seed", default=12345, help="seed for variance testing", type=int)

        # General parameters
        parser.add_argument("--save_count", default=100, help="Number of ckpts for saving results and model", type=int)
        parser.add_argument("--optim", default='adam', help="Optimizer type", choices=['adam', 'sgd', 'rmsprop'])
        parser.add_argument("--log_output", default='term_file', help="Log all the print outputs",
                            choices=['term_file', 'term', 'file'])
        parser.add_argument("--debug", default=True, type=self.str2bool, help="Debug mode on/off")
        parser.add_argument("--restore", default=False, type=self.str2bool, help="Retrain flag")
        parser.add_argument("--save_model", default=True, type=self.str2bool, help="flag to save model ckpts")
        parser.add_argument("--gpu", default=0, help="GPU BUS ID ", type=int)

        # Book-keeping parameters
        now = datetime.now()
        timestamp = str(now.month) + '|' + str(now.day) + '|' + str(now.hour) + ':' + str(now.minute) + ':' + str(
            now.second)
        parser.add_argument("--timestamp", default=timestamp, help="Timestamp to prefix experiment dumps")
        parser.add_argument("--folder_suffix", default='Default', help="folder name suffix")
        parser.add_argument("--experiment", default='Test_run', help="Name of the experiment")
        parser.add_argument("--domain", default='land', help="land or air")

        parser.add_argument("--model_lr", default=1e-3, help="Learning rate of actor", type=float)
        parser.add_argument("--controller_lr", default=1e-3, help="Learning rate of critic/baseline", type=float)
        parser.add_argument("--noise_var", default=1e-2, help="Variance of noise in dynamics", type=float)

        parser.add_argument("--valid_fraction", default=0.2, help="Fraction of data for validation", type=float)
        parser.add_argument("--hidden_nodes", default=16, help="hidden nodes in NN", type=int)
        parser.add_argument("--max_steps", default=50, help="maximum number of steps in trajectory", type=int)
        parser.add_argument("--max_episodes", default=1000, help="maximum episodes for training controller", type=int)
        parser.add_argument("--batch_size", default=32, help="Batch size for estimating system dynamics", type=int)
        parser.add_argument("--max_model_iterations", default=1000, help="Max epochs for learning dynamics", type=int)
        parser.add_argument("--true_model", default=False, type=self.str2bool, help="Use true or estimated  model?")

        self.parser = parser

    def str2bool(self, text):
        if text == 'True':
            arg = True
        elif text == 'False':
            arg = False
        else:
            raise argparse.ArgumentTypeError('Boolean value expected.')
        return arg

    def get_parser(self):
        return self.parser
