from torch.optim import Adam, NAdam, RAdam, AdamW, Adamax


def adam(parameters, lr, betas, eps, weight_decay):
    '''
    Return Adam optimizer object.
    https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam

    Parameterts
    -----------
    parameters : Any
        model.parameters()
    lr : float
        learning rate
    betas : tuple(float, float)
        betas
    eps : float
        eps
    weight_decay : float
        l2 regularization
    '''
    return Adam(parameters, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)


def nadam(parameters, lr, betas, eps, weight_decay):
    '''
    Return NAdam optimizer object.
    https://pytorch.org/docs/stable/generated/torch.optim.NAdam.html#torch.optim.NAdam

    Parameterts
    -----------
    parameters : Any
        model.parameters().
    lr : float
        learning rate.
    betas : tuple(float, float)
        betas.
    eps : float
        eps.
    weight_decay : float
        l2 regularization.
    '''
    return NAdam(parameters, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)


def radam(parameters, lr, betas, eps, weight_decay):
    '''
    Return RAdam optimizer object.
    https://pytorch.org/docs/stable/generated/torch.optim.RAdam.html#torch.optim.RAdam

    Parameterts
    -----------
    parameters : Any
        model.parameters().
    lr : float
        learning rate.
    betas : tuple(float, float)
        betas.
    eps : float
        eps.
    weight_decay : float
        l2 regularization.
    '''
    return RAdam(parameters, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)


def adamw(parameters, lr, betas, eps, weight_decay):
    '''
    Return AdamW optimizer object.
    https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html#torch.optim.AdamW

    Parameterts
    -----------
    parameters : Any
        model.parameters().
    lr : float
        learning rate.
    betas : tuple(float, float)
        betas.
    eps : float
        eps.
    weight_decay : float
        l2 regularization.
    '''
    return AdamW(parameters, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)


def adamax(parameters, lr, betas, eps, weight_decay):
    '''
    Return Adamax optimizer object.
    https://pytorch.org/docs/stable/generated/torch.optim.Adamax.html#torch.optim.Adamax

    Parameterts
    -----------
    parameters : Any
        model.parameters().
    lr : float
        learning rate.
    betas : tuple(float, float)
        betas.
    eps : float
        eps.
    weight_decay : float
        l2 regularization.
    '''
    return Adamax(parameters, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
