import torch


def packed_to_padded(seq, target_packed):
    """Converts a sequence of packed outputs into a padded tensor

    Arguments:
        seq {list} -- List of lists of length T (longest sequence) where each
            sub-list contains output tokens for relevant samples.
        E.g. [len(s) for s in seq] == [8, 8, 8, 4, 3, 1] if batch has 8 samples
        longest sequence has length 6 and only 3/8 samples have length 3.
        target_packed {list} -- Packed target sequence
    Return:
        torch.Tensor: Shape bs x T (padded with 0)

    NOTE:
        Assumes that padding index is 0 and stop_index is 3
    """
    T = len(seq)
    batch_size = len(seq[0])
    padded = torch.zeros(batch_size, T)

    stopped_idx = []
    target_packed += [torch.Tensor()]
    # Loop over tokens per time step
    for t in range(T):
        seq_lst = seq[t].tolist()
        tg_lst = target_packed[t - 1].tolist()
        # Insert Padding token where necessary
        [seq_lst.insert(idx, 0) for idx in sorted(stopped_idx, reverse=False)]
        padded[:, t] = torch.Tensor(seq_lst).long()

        stop_idx = list(filter(lambda x: tg_lst[x] == 3, range(len(tg_lst))))
        stopped_idx += stop_idx

    return padded