import torch


def concat_hidden(beam, nlayers, m_per_layer=2):
    """
    takes a beam of Candidates and a number of layers and makes new concatenated layers for batching efficiency
    :param beam:
    :param nlayers:
    :param m_per_layer: matrices per later. Defaults to 2 for LSTMs
    :return: list of tuples, one tuple per layer, that are concatenation of all layers belonging
    to candidates
    """
    new_hidden = []
    for l in range(nlayers):
        # need to add an additional dimension before concatenation else get (1, 7500) instead of (1, 5, 1500) with beam of 5 and hidden layers 1500
        new_layer = tuple([torch.cat([cand.hidden[l][i].unsqueeze(1) for cand in beam], dim=1)
                           for i in range(m_per_layer)])
        new_hidden.append(new_layer)
    return new_hidden