import torch


def decode_to_token(x, idx_to_tok, args, mask=None):
    """
    Given a batch of embeddings, returns string tokens. Masked items will be
    removed from output. If no mask is provided, it will be created on-the-fly
    and returned.

    Args:
        x: (Tensor) size [batch, seq length, emb size]
        idx_to_tok: (list) return the string token given index where the index
                    matches to the embedding index
        mask: (Tensor) same shape as `x`
    """
    # Create mask
    if mask is None:
        mask = torch.where(x > 0, torch.tensor([1.]).to(args.device),
                                            torch.tensor([0.]).to(args.device))
    x_tokens = []
    # Loop sequences
    for j in range(x.shape[0]):
        seq_tokens = []
        for i, word in enumerate(x[j]):
            if float(mask[j,i]) > 0:
                seq_tokens.append(idx_to_tok[int(x[j,i])])
        sentence = ' '.join(seq_tokens[:])
        x_tokens.append(sentence)
    return x_tokens, mask