# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_TorchUtils.ipynb.

# %% auto 0
__all__ = ['device_by_name', 'DatasetCat']

# %% ../nbs/01_TorchUtils.ipynb 3
from torch.utils.data import Dataset
import torch

# %% ../nbs/01_TorchUtils.ipynb 4
def device_by_name(name: str) -> torch.device:
    ''' Return reference to cuda device by using Part of it's name

        Args:
            name: part of the cuda device name (shuuld be distinct)

        Return:
            Reference to cuda device

        Updated: Yuval 12/10/19
    '''
    assert torch.cuda.is_available(), "No cuda device"
    device = None
    for i in range(torch.cuda.device_count()):
        dv = torch.device("cuda:{}".format(i))
        if name in torch.cuda.get_device_name(dv):
            device = dv
            break
    assert device, "device {} not found".format(name)
    return device

# %% ../nbs/01_TorchUtils.ipynb 10
class DatasetCat(Dataset):
    '''
    Concatenate datasets for Pytorch dataloader
    
    The normal pytorch implementation does it only for raws. this is a "column" implementation
    
    Arges:
        datasets: list of datasets, of the same length
        
    Updated: Yuval 12/10/2019
    '''

    def __init__(self, *datasets):
        '''
        Args: datasets - an iterable containing the datasets
        '''
        super(DatasetCat, self).__init__()
        self.datasets=datasets
        assert len(self.datasets)>0
        for dataset in datasets:
            assert len(self.datasets[0])==len(dataset),"Datasets length should be equal"

    def __len__(self):
        return len(self.datasets[0])

    def __getitem__(self, idx):
        outputs = tuple(dataset.__getitem__(idx) for i in self.datasets for dataset in (i if isinstance(i, tuple) else (i,)))
        return tuple(output for i in outputs for output in (i if isinstance(i, tuple) else (i,)))
    
