import torch


def init_optimizers(args, params, alpha, beta):
    """
    Initialize optimizers for the quantization network model.

    Args:
        args: The args object obtained with argparse.
            The following requirements should be met:
            args.lr is a valid floating point number (0~1).
            args.weight_decay is a valid floating point number (0~1).
            args.momentum is a valid floating point number (0~1).
        params: All the parameters to optimize.
        alpha: All the alpha values to optimize.
        beta: All the beta values to optimize.

    Returns:
        A tuple of 3 optimizer objects, one for network weights, one for alphas and one for betas.
    """
    optimizer_alpha = None
    optimizer_beta = None
    if args.adam:
        print('The optimizer is Adam !!!')
        optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
        if args.qw:
            optimizer_alpha = torch.optim.Adam(alpha, lr=args.lr)
            optimizer_beta = torch.optim.Adam(beta, lr=args.lr)
    else:
        print('The optimizer is SGD !!!')
        optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
        if args.qw:
            optimizer_alpha = torch.optim.SGD(alpha, lr=args.lr, momentum=args.momentum)
            optimizer_beta = torch.optim.SGD(beta, lr=args.lr, momentum=args.momentum)
    return optimizer, optimizer_alpha, optimizer_beta