import torch


def agg_logs(outputs):
    """
    Aggregate a list of dicts into a single (may have sub dicts but all array are aggregated)

    outputs = [
        {'val_loss': 0.7206,
            'log': {'val_loss': 0.7206, 'val_loss_p': 0.7206,}},
        {'val_loss': 0.7047,
            'log': {'val_loss': 0.7047, 'val_loss_p': 0.7047}},
    ]
    -> {'agg_val_loss': 0.7126500010490417, 'log': {'agg_val_loss': 0.7126500010490417, 'agg_val_loss_p': 0.7126500010490417, 'agg_val_loss_kl': 2.6101499770447845e-06, 'agg_val_loss_mse': 0.17669999599456787}}
    
    """
    if isinstance(outputs, dict):
        outputs = [outputs]
    
    aggs = {}
    if len(outputs) > 0:
        for j in outputs[0].keys():
            if isinstance(outputs[0][j], dict):
                # Take mean of sub dicts
                keys = outputs[0][j].keys()
                aggs[j] = {
                    'agg_'+k: torch.stack([x[j][k] for x in outputs if k in x[j]])
                    .mean()
                    .cpu()
                    .item()
                    for k in keys
                }
            else:
                # Take mean of numbers
                aggs['agg_'+j] = (
                    torch.stack([x[j] for x in outputs if j in x]).mean().cpu().item()
                )
    return aggs