import torch


def init_schedulers(args, optimizer, optimizer_alpha, optimizer_beta):
    """
    Initialize LR schedulers.

    Args:
        args: The args object obtained with argparse.
            The following requirements should be met:
            args.qw is a boolean value indicating whether the model weights should be quantized.
        optimizer: The optimizer for the model.
        optimizer_alpha: The optimizer for the alpha variables in weight quantization. Can be None.
        optimizer_beta: The optimizer for the beta variables in weight quantization. Can be None.

    Returns:
        A tuple of 3 elements:
            lr_scheduler: The LR scheduler for the model's optimizer.
            lr_scheduler_alpha: The LR scheduler for the alpha optimizer.
            lr_scheduler_beta: The LR scheduler for the beta optimizer.
    """
    lr_scheduler_alpha = None
    lr_scheduler_beta = None
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.decay_steps, last_epoch=-1)
    if args.qw:
        lr_scheduler_alpha = torch.optim.lr_scheduler.MultiStepLR(optimizer_alpha, milestones=args.decay_steps,
                                                                  last_epoch=-1)
        lr_scheduler_beta = torch.optim.lr_scheduler.MultiStepLR(optimizer_beta, milestones=args.decay_steps,
                                                                 last_epoch=-1)

    return lr_scheduler, lr_scheduler_alpha, lr_scheduler_beta