import torch


def cdist(a, b, metric='euclidean'):
    """Similar to scipy.spatial's cdist, but symbolic.
    The currently supported metrics can be listed as `cdist.supported_metrics` and are:
        - 'euclidean', although with a fudge-factor epsilon.
        - 'sqeuclidean', the squared euclidean.
        - 'cityblock', the manhattan or L1 distance.
    Args:
        a: The left-hand side, shaped ([*,] F, B1).  <- Not that dimension ordering is different from torch.cdist
        b: The right-hand side, shaped ([*,], F, B2).
        metric (string): Which distance metric to use, see notes.
    Returns:
        The matrix of all pairwise distances between all vectors in `a` and in
        `b`, will be of shape (B1, B2).
    Note:
        When a square root is taken (such as in the Euclidean case), a small
        epsilon is added because the gradient of the square-root at zero is
        undefined. Thus, it will never return exact zero in these cases.

    Taken from Predator source code, which was modified from D3Feat.
    """
    if metric == 'sqeuclidean':
        diffs = a[..., :, None] - b[..., None, :]
        return torch.sum(diffs ** 2, dim=-3)
    elif metric == 'euclidean':
        diffs = a[..., :, None] - b[..., None, :]
        return torch.sqrt(torch.sum(diffs ** 2, dim=-3) + 1e-12)
    elif metric == 'cityblock':
        diffs = a[..., :, None] - b[..., None, :]
        return torch.sum(torch.abs(diffs), dim=-3)
    elif metric == 'cosine':
        numer = a.transpose(-1, -2) @ b
        denom = torch.clamp_min(
            torch.norm(a, dim=-2)[..., :, None] * torch.norm(b, dim=-2)[..., None, :],
            1e-8)
        dist = 1 - numer / denom
        return dist
    else:
        raise NotImplementedError(
            'The following metric is not implemented by `cdist` yet: {}'.format(metric))