import types
import torch


def get_optimizer(params, lr, momentum=0, weight_decay=0, use_adam=False,
                  adam_beta1=0.9, use_rmsprop=False, use_adadelta=False,
                  use_adagrad=False, pgroup_ids=None):
    """Create an optimizer instance for the given set of parameters. Default
    optimizer is :class:`torch.optim.SGD`.

    Args:
        params (list): The parameters passed to the optimizer.
        lr: Learning rate.
        momentum (optional): Momentum (only applicable to
            :class:`torch.optim.SGD` and :class:`torch.optim.RMSprop`.
        weight_decay (optional): L2 penalty.
        use_adam: Use :class:`torch.optim.Adam` optimizer.
        adam_beta1: First parameter in the `betas` tuple that is passed to the
            optimizer :class:`torch.optim.Adam`:
            :code:`betas=(adam_beta1, 0.999)`.
        use_rmsprop: Use :class:`torch.optim.RMSprop` optimizer.
        use_adadelta: Use :class:`torch.optim.Adadelta` optimizer.
        use_adagrad: Use :class:`torch.optim.Adagrad` optimizer.
        pgroup_ids (list, optional): If passed, a list of integers of the same
            length as params is expected. In this case, each integer states to
            which parameter group the corresponding parameter in ``params``
            shall belong. Parameter groups may have different optimizer
            settings. Therefore, options like ``lr``, ``momentum``,
            ``weight_decay``, ``adam_beta1`` may be lists in this case that have
            a length corresponding to the number of parameter groups.

    Returns:
        Optimizer instance.
    """
    use_sgd = not use_adam and not use_rmsprop and not use_adadelta \
        and not use_adagrad

    if isinstance(params, types.GeneratorType):
        params = list(params)

    # Transform list of parameter tensors `params` into list of dictionaries.
    if pgroup_ids is None:
        pgroup_ids = [0] * len(params)
        npgroups = 1
    else:
        assert len(pgroup_ids) == len(params)
        npgroups = max(pgroup_ids) + 1

    plist = params
    params = []

    # Initialize parameter group dictionaries.
    for i in range(npgroups):
        pdict = {}
        pdict['params'] = []
        pdict['lr'] = lr[i] if isinstance(lr, (list, tuple)) else lr
        pdict['weight_decay'] = weight_decay[i] \
            if isinstance(weight_decay, (list, tuple)) else weight_decay
        if use_adam:
            ab1 = adam_beta1[i] if isinstance(adam_beta1, (list, tuple)) \
                else adam_beta1
            pdict['betas'] = [ab1, 0.999]
        if use_sgd or use_rmsprop:
            pdict['momentum'] = momentum[i] \
                if isinstance(momentum, (list, tuple)) else momentum
        params.append(pdict)

    # Fill parameter groups.
    for pgid, p in zip(pgroup_ids, plist):
        params[pgid]['params'].append(p)


    if use_adam:
        optimizer = torch.optim.Adam(params)
    elif use_rmsprop:
        optimizer = torch.optim.RMSprop(params)
    elif use_adadelta:
        optimizer = torch.optim.Adadelta(params)
    elif use_adagrad:
        optimizer = torch.optim.Adagrad(params)
    else:
        assert use_sgd
        optimizer = torch.optim.SGD(params)

    return optimizer