import torch
from typing import Optional
from typing import Sequence


def correct(outputs: torch.Tensor, labels: torch.Tensor,
            weights: Optional[torch.Tensor] = None,
            top: Sequence[int] = (1,)) -> dict[int, float]:
    """
    Args:
        outputs: torch.Tensor, shape [N, num_classes],
            either logits (pre-softmax) or probabilities
        labels: torch.Tensor, shape [N]
        weights: optional torch.Tensor, shape [N]
        top: tuple of int, list of values of k for calculating top-K accuracy

    Returns: dict, maps k to (weighted) # of correct predictions @ each k
    """
    with torch.no_grad():
        # preds and labels both have shape [N, k]
        _, preds = outputs.topk(k=max(top), dim=1, largest=True, sorted=True)
        labels = labels.view(-1, 1).expand_as(preds)

        corrects = preds.eq(labels).cumsum(dim=1)  # shape [N, k]
        if weights is None:
            corrects = corrects.sum(dim=0)  # shape [k]
        else:
            corrects = weights.matmul(corrects.to(weights.dtype))  # shape [k]
        tops = {k: corrects[k - 1].item() for k in top}
    return tops