from typing import Tuple
from typing import Union
from typing import List
import torch


def device_to_name_ids(device: str) -> Tuple[str, Union[None, List[int]]]:
    """
    Split a device string into a device and ids

    :param device: the device string to push to; ex: cpu, cuda, cuda:0,1
    :return: a tuple containing the device string and devices
    """
    split = device.split(":")
    name = split[0]

    if name == "cpu":
        return name, None

    if name != "cuda" or not torch.cuda.is_available():
        raise ValueError("{} device not available on this system".format(name))

    if len(split) < 2:
        return name, None

    ids = [int(id_) for id_ in split[1].split(",")]
    count = torch.cuda.device_count()

    for id_ in ids:
        if id_ >= count:
            raise ValueError("{} device id not available on this system".format(id_))

    if len(ids) == 1:
        return "{}:{}".format(name, ids[0]), None

    return name, ids