import torch


def choose_forward_optimizer(args, net):
    """
    Return the wished optimizer (based on inputs from args).
    Args:
        args: cli
        net: neural network
    Returns: optimizer
    """

    if args.network_type in ('BP'):
        if args.shallow_training:
            print('Shallow training')
            forward_params = net.layers[-1].parameters()
        elif args.only_train_first_layer:
            print('Only training first layer')
            forward_params = net.layers[0].parameters()
        else:
            forward_params = net.parameters()
    else:
        if args.only_train_first_layer:
            print('Only training first layer')
            forward_params = net.get_forward_parameter_list_first_layer()
        else:
            forward_params = net.get_forward_parameter_list()

    if args.optimizer == 'SGD':
        print('Using SGD optimizer')
        forward_optimizer = torch.optim.SGD(forward_params,
                                            lr=args.lr, momentum=args.momentum,
                                            weight_decay=args.forward_wd)
    elif args.optimizer == 'RMSprop':
        print('Using RMSprop optimizer')

        forward_optimizer = torch.optim.RMSprop(
            forward_params,
            lr=args.lr,
            momentum=args.momentum,
            alpha=0.95,
            eps=0.03,
            weight_decay=args.forward_wd,
            centered=True
        )

    elif args.optimizer == 'Adam':
        print('Using Adam optimizer.')

        forward_optimizer = torch.optim.Adam(
            forward_params,
            lr=args.lr,
            betas=(args.beta1, args.beta2),
            eps=args.epsilon,
            weight_decay=args.forward_wd
        )

    else:
        raise ValueError('Provided optimizer "{}" is not supported'.format(
            args.optimizer
        ))

    return forward_optimizer