import typing
import torch


def join_batches(batches: typing.List[typing.Dict]):
    """ Joins encoded batches to build a bigger batch

    Args:
        batches:
            List of encoded batches. Each batch is a dict with keys
            (events, timestamps, images, augmentation_params).

    Returns:
        A joined batch as a dict with keys
        (events, timestamps, images, augmentation_params).
    """

    if len(batches) == 0:
        return {'events': {'x': torch.tensor([], dtype=torch.short),
                           'y': torch.tensor([], dtype=torch.short),
                           'timestamp': torch.tensor([], dtype=torch.float32),
                           'polarity': torch.tensor([], dtype=torch.bool),
                           'events_per_element':
                           torch.tensor([], dtype=torch.short),
                           'elements_per_sample':
                           torch.tensor([], dtype=torch.short)},
                'timestamps': torch.tensor([], dtype=torch.float32),
                'images': torch.tensor([], dtype=torch.uint8),
                'augmentation_params': {}}
    if len(batches) == 1:
        return batches[0]
    result = {}
    for k in batches[0].keys():
        if isinstance(batches[0][k], dict):
            result[k] = {}
            for sk in batches[0][k].keys():
                result[k][sk] = torch.cat([el[k][sk] for el in batches])
        elif batches[0][k] is None:
            assert k == 'augmentation_params'
            assert all([el[k] is None for el in batches])
            result[k] = None
        else:
            assert isinstance(batches[0][k], torch.Tensor)
            result[k] = torch.cat([el[k] for el in batches])
    return result