import torch


def topk_accuracy(k, logits: torch.Tensor, labels: torch.Tensor, unk_id=None, pad_id=None):
    """
    Averaged per sample:
        1, if one of the k predictions with highest confidence was the correct one for all sub tokens
        0, else
    if `unk_id` is given, the best non <unk> prediction will be used and <unk> tokens in the labels will be
    completely ignored, i.e., it does not matter what the 5 predictions with highest score at this position were.
    :param logits: batch_size x num_predict x num_sub_tokens x vocab_size
    :param labels: batch_size x num_predict x num_sub_tokens
    :param k: The k highest logits will be considered as predictions and for every sub token when any of the k
    highest valued predictions was the correct one, it will count as an accurate prediction
    :param unk_id: ID of the <unk> token
    """
    assert unk_id is None or pad_id is not None, "When unk_id is given, pad_id must be given as well"
    topk_pred = logits.topk(k, -1)
    # Accept if any of the top k predictions is the label
    topk_correct = (topk_pred.indices == labels.unsqueeze(-1).expand((-1, -1, -1, k)))
    topk_correct = topk_correct.any(-1)
    if unk_id is None:
        topk_correct = topk_correct.all(-1)  # Only accept if for all sub tokens a top k prediction was correct
    if unk_id is not None:
        idx_unk_labels = (labels == unk_id)
        topk_correct[idx_unk_labels] = True
        topk_correct = topk_correct.all(-1)  # Only accept if for all sub tokens a top k prediction was correct
        idx_label_all_unk = ((labels == unk_id) | (labels == pad_id)).all(-1) & ~(labels == pad_id).all(-1)
        topk_correct = topk_correct[~idx_label_all_unk]  # Completely ignore if labels only
        # consists of <unk>

    return topk_correct.float().mean().item()