import torch


def compute_lid(x, x_train, k, exclude_self=False):
    """
    Calculate LID using the estimation from [1]

    [1] Ma et al., "Characterizing Adversarial Subspaces Using
        Local Intrinsic Dimensionality," ICLR 2018.
    """

    with torch.no_grad():
        x = x.view((x.size(0), -1))
        x_train = x_train.view((x_train.size(0), -1))
        lid = torch.zeros((x.size(0), ))

        for i, x_cur in enumerate(x):
            dist = (x_cur.view(1, -1) - x_train).norm(2, 1)
            # `largest` should be True when using cosine distance
            if exclude_self:
                topk_dist = dist.topk(k + 1, largest=False)[0][1:]
            else:
                topk_dist = dist.topk(k, largest=False)[0]
            mean_log = torch.log(topk_dist / topk_dist[-1]).mean()
            lid[i] = -1 / mean_log
        return lid