from argparse import ArgumentParser

class HyperParameters():
    def parse(self, unknown_arg_ok=False):
        parser = ArgumentParser()

        # Generic learning parameters
        parser.add_argument('-i', '--iterations', help='Number of training iterations', default=4.5e4, type=int)
        parser.add_argument('-b', '--batch_size', help='Batch size', default=12, type=int)
        parser.add_argument('--lr', help='Initial learning rate', default=2.25e-4, type=float)
        parser.add_argument('--steps', help='Iteration at which learning rate is decayed by gamma', default=[22500, 37500], type=int, nargs='*')
        parser.add_argument('--gamma', help='Gamma used in learning rate decay', default=0.1, type=float)
        parser.add_argument('--weight_decay', help='Weight decay', default=1e-4, type=float)

        # same decay applied to discriminator
        parser.add_argument('--load', help='Path to pretrained model if available')

        parser.add_argument('--ce_weight', help='Weight of CE loss function for each iteration',
            nargs=6, default=[0.0, 1.0, 0.5, 1.0, 1.0, 0.5], type=float)
        parser.add_argument('--l1_weight', help='Weight of L1 loss function for each iteration',
            nargs=6, default=[1.0, 0.0, 0.25, 0.0, 0.0, 0.25], type=float)
        parser.add_argument('--l2_weight', help='Weight of L2 loss function for each iteration',
            nargs=6, default=[1.0, 0.0, 0.25, 0.0, 0.0, 0.25], type=float)
        parser.add_argument('--grad_weight', help='Weight of the gradient loss', default=5, type=float)

        # Logging information, this one is positional and mandatory
        parser.add_argument('id', help='Experiment UNIQUE id, use NULL to disable logging to tensorboard')

        if unknown_arg_ok:
            args, _ = parser.parse_known_args()
            self.args = vars(args)
        else:
            self.args = vars(parser.parse_args())

    def __getitem__(self, key):
        return self.args[key]

    def __str__(self):
        return str(self.args)

