import torch


def extract_torchmeta_task(cs, class_ids):
    """ Extracts a single "episode" (ie, task) from a ClassSplitter object, in the
        form of a dataset, and appends variables needed by DatasetDistance computation.

        Arguments:
            cs (torchmeta.transforms.ClassSplitter): the ClassSplitter where to extract data from
            class_ids (tuple): indices of classes to be selected by Splitter

        Returns:
            ds_train (Dataset): train dataset
            ds_test (Dataset): test dataset

    """
    ds = cs[class_ids]
    ds_train, ds_test = ds['train'], ds['test']

    for ds in [ds_train, ds_test]:
        ds.targets = torch.tensor([ds[i][1] for i in range(len(ds))])
        ds.classes = [p[-1] for i,p in enumerate(cs.dataset._labels) if i in class_ids]
    return ds_train, ds_test