from typing import Tuple
from typing import Dict
import torch


def classification_collate(batch: Tuple[Dict]) -> Dict:
    """Collate function for classification model

    :param batch: A tuple of dicts of processed signals, labels and items
    :type batch: Tuple[Dict]
    :returns: A dict containing:
        1) tensor, batch of processed signals
        2) tensor, batch of corresponding labels
        3) list, items corresponding to each instance
    """
    signals = []
    labels = []
    items = []

    for data_point in batch:
        signals.append(data_point['signal'])
        labels.append(data_point['label'])
        items.append(data_point['item'])

    collated_batch = {
        'signals': torch.stack(signals),
        'labels': torch.Tensor(labels),
        'items': items
    }

    return collated_batch