from torch import optim
from typing import Dict
from typing import Any
import torch


def get_torch_optimizer(
    name: str, hparams: Dict[str, Any], parameters
) -> torch.optim.Optimizer:
    """Construct a PyTorch optimizer specified by :obj:`name` and :obj:`hparams`."""

    if name == "Adadelta":
        optimizer = optim.Adadelta
    elif name == "Adagrad":
        optimizer = optim.Adagrad
    elif name == "Adam":
        optimizer = optim.Adam
    elif name == "AdamW":
        optimizer = optim.AdamW
    elif name == "SparseAdam":
        optimizer = optim.SparseAdam
    elif name == "Adamax":
        optimizer = optim.Adamax
    elif name == "ASGD":
        optimizer = optim.ASGD
    elif name == "LBFGS":
        optimizer = optim.LBFGS
    elif name == "RMSprop":
        optimizer = optim.RMSprop
    elif name == "Rprop":
        optimizer = optim.Rprop
    elif name == "SGD":
        optimizer = optim.SGD
    else:
        raise ValueError(f"Invalid optimizer name: {name}")

    try:
        return optimizer(parameters, **hparams)

    except TypeError:
        raise Exception(
            f"Invalid parameter in hparams: {hparams}"
            f" for optimizer {name}.\nSee PyTorch docs."
        )