import numpy as np

from utils.utils import normalize


def get_summary_dict(batch, run, visualization_keys=None, *others):
    if visualization_keys is None:
        visualization_keys = ['reconstruction', 'L1']
    run = dict(filter(lambda x: x[1] is not None, run.items()))
    visuals = np.asarray([
        255 * np.hstack([
            normalize(batch[i]),
            *[normalize(run[key][i]) for key in visualization_keys],
            *[normalize(element[i]) for element in others]
        ]) for i in range(len(batch))]
    )
    scalars = dict(filter(lambda x: not (type(x[1]) == float and x[1] != x[1]) and x[1].ndim == 0, run.items()))
    return scalars, visuals