import torch


def compute_rank(predictions, targets):
    """Compute the rank (between 1 and n) of of the true target in ordered predictions

    Example:
        >>> import torch
        >>> compute_rank(torch.tensor([[.1, .7, 0., 0., .2, 0., 0.],
        ...                            [.1, .7, 0., 0., .2, 0., 0.],
        ...                            [.7, .2, .1, 0., 0., 0., 0.]]),
        ...              torch.tensor([4, 1, 3]))
        tensor([2, 1, 5])

    Args:
        predictions (torch.Tensor): [n_pred, n_node]
        targets (torch.Tensor): [n_pred]
    """
    n_pred = predictions.shape[0]
    range_ = torch.arange(n_pred, device=predictions.device, dtype=torch.long)

    proba_targets = predictions[range_, targets]
    target_rank_upper = (proba_targets.unsqueeze(1) < predictions).long().sum(dim=1) + 1
    target_rank_lower = (proba_targets.unsqueeze(1) <= predictions).long().sum(dim=1)

    # break tighs evenly by taking the mean rank
    target_rank = (target_rank_upper + target_rank_lower) / 2
    return target_rank