import torch
from typing import get_args


def create_optim(enc, parameters, **kwargs):
    """Initialise optimisers.

    Args:
      enc (string): type of optimiser - either 'SGD' or 'Adam'.
      parameters (iterable): parameters to be optimised.

    Returns:
      An instance of torch.optim.

    Raises:
      ValueError if enc is not either of 'SGD' or 'Adam'.

    """
    if enc == 'SGD':
        optim = torch.optim.SGD
    elif enc == 'Adam':
        optim = torch.optim.Adam
    else:
        raise ValueError("Optim {} is not supported. "
                         "Only supports 'SGD' and 'Adam' for now.".format(enc))
    args = get_args(optim)
    kwargs = {key : kwargs[key] for key in args if key in kwargs}
    return optim(parameters, **kwargs)