import torch


def pad_tensor(adj_nodes_list, mask=False):
    """Function pads the neighbourhood nodes before passing through the 
    aggregator.

    Args:
        adj_nodes_list (list): the list of node neighbours
        mask (bool, optional): if true, create a tensor with 1s and 0s for masking. Defaults to False.

    Returns:
        tuple: one of two tensors containing the padded tensor and mask (if true)
    """
    max_len = max([len(adj_nodes) for adj_nodes in adj_nodes_list])

    padded_nodes = []
    _mask = []
    for adj_nodes in adj_nodes_list:
        x = list(adj_nodes)
        
        x += [0] * (max_len - len(adj_nodes))
        padded_nodes.append(x)
        _mask.append([1] * len(adj_nodes) + [0] * (max_len - len(adj_nodes)))

    if not mask:
        return torch.tensor(padded_nodes)

    # returning the mask as well
    return torch.tensor(padded_nodes), torch.tensor(_mask)