import torch


def collate_fn(task_list):
	"""Custom collate function to batch tasks together into single tensors.

	Args:
	    task_list: a list of tasks, where each task is a dictionary as returned by the
                       __getitem__ method of the dataset object

	Returns:
	    batch: a dictionary with keys 'cond', 'pool', 'pool_labels', 'query', 'query_labels'.
		   The value of 'cond' is an array of shape (batch_size), the value of 'pool' is a
                   tensor of shape (batch_size, pool_size, embedding_size), the value of 'pool_labels'
                   is a tensor of shape (batch_size, pool_size), the value of 'query' is a tensor of
                   shape (batch_size, query_set_size, embedding_size), and the value of 'query_labels'
                   is a tensor of shape (batch_size, query_set_size)
	"""
	cond = [s['cond'] for s in task_list]
	batch = {'cond': cond, 'pool': None, 'pool_labels': None, 'query': None, 'query_labels': None}

	keys = ['pool', 'pool_labels', 'query', 'query_labels']
	for key in keys:
		tensor_list = [s[key] for s in task_list]
		stacked = torch.stack(tensor_list, dim=0)
		batch[key] = stacked

	batch['pool_labels'] = batch['pool_labels'].long()
	batch['query_labels'] = batch['query_labels'].long()
	return batch