from typing import OrderedDict
import torch


def collate_metrics(output_data_batch, name='depth'):
    """
    Collate epoch output to produce average metrics

    Parameters
    ----------
    output_data_batch : list
        List of outputs for each batch
    name : str
        Name of the task for the metric

    Returns
    -------
    metrics_data : list
        List of collated metrics
    """
    # If there is only one dataset, wrap in a list
    if isinstance(output_data_batch[0], dict):
        output_data_batch = [output_data_batch]
    # Calculate the mean of all metrics
    metrics_data = []
    # For all datasets
    for i, output_batch in enumerate(output_data_batch):
        metrics = OrderedDict()
        # For all keys (assume they are the same for all batches)
        for key, val in output_batch[0].items():
            if key.startswith(name):
                metrics[key] = torch.stack([output[key] for output in output_batch], 0)
                metrics[key] = torch.mean(metrics[key], 0)
        metrics_data.append(metrics)
    # Return metrics data
    return metrics_data