import torch


def _label_propagate(features, y_l, threshold=0.95):
    """Perform random walk based label propagation with similarity graph.

    Arguments:
        features: features of size (N, D), where N is the number of superpixels
            and D is the dimension of input features
        y_l: label tensor of size (N, C), where C is the number of classes
        threshold: similarity threshold for label propagation

    Returns:
        pseudo_labels: propagated label tensor of size (N, C)
    """

    # disable gradient computation
    features = features.detach()
    y_l = y_l.detach()

    # number of labeled and unlabeled samples
    n_l = y_l.size(0)
    n_u = features.size(0) - n_l

    # feature affinity matrix
    W = torch.exp(-torch.einsum('ijk,ijk->ij',
                                features - features.unsqueeze(1),
                                features - features.unsqueeze(1)))

    # sub-matrix of W containing similarities between labeled and unlabeled samples
    W_ul = W[n_l:, :n_l]

    # max_similarities is the maximum similarity for each unlabeled sample
    # src_indexes is the respective labeled sample index
    max_similarities, src_indexes = W_ul.max(dim=1)

    # initialize y_u with zeros
    y_u = torch.zeros(n_u, y_l.size(1)).to(y_l.device)

    # only propagate labels if maximum similarity is above the threshold
    propagated_samples = max_similarities > threshold
    y_u[propagated_samples] = y_l[src_indexes[propagated_samples]]

    return y_u