from typing import Union
from typing import Sequence
import torch
from typing import Mapping


def to_device(inp: Union[Sequence[torch.Tensor], torch.Tensor, Mapping[str, torch.Tensor]],
              device: Union[torch.device, str], detach: bool = False,
              **kwargs) -> \
        Union[Sequence[torch.Tensor], torch.Tensor, Mapping[str, torch.Tensor]]:
    """
    Push tensor or sequence of tensors to device

    Args:
        inp (Union[Sequence[torch.Tensor], torch.Tensor]): tensor or sequence of tensors
        device (Union[torch.device, str]): target device
        detach: detach tensor before moving it to new device
        **kwargs: keyword arguments passed to `to` function

    Returns:
        Union[Sequence[torch.Tensor], torch.Tensor]: tensor or seq. of tenors at target device
    """
    if isinstance(inp, torch.Tensor):
        if detach:
            return inp.detach().to(device=device, **kwargs)
        else:
            return inp.to(device=device, **kwargs)
    elif isinstance(inp, Sequence):
        old_type = type(inp)
        return old_type([to_device(i, device=device, detach=detach, **kwargs)
                         for i in inp])
    elif isinstance(inp, Mapping):
        old_type = type(inp)
        return old_type({key: to_device(item, device=device, detach=detach, **kwargs)
                         for key, item in inp.items()})
    else:
        return inp