import torch


def get_matching(samples, labels, topk=1, max_err=False):
    """Get mean or max matching over images."""
    assert topk == 1

    # has shape (num_f, num_eps, num_images, num_samples)
    # NOTE: labels must be broadcastable to that ...
    idxs_top = samples.topk(topk, dim=-1)[1].squeeze(dim=-1)

    # compare matching labels
    matching = (idxs_top == labels).float()

    if max_err:
        # raise NotImplementedError
        num_samples = matching.shape[-2]
        percentile = 0.01
        # reduced = torch.min(matching, dim=-2)[0]
        reduced = torch.median(matching, dim=-2)[0]
        reduced = torch.sort(matching, dim=-2)[0][
            :, :, int(percentile * num_samples)
        ]
    else:
        # mean over images (i.e. "Test Accuracy" over noisy images )
        reduced = torch.mean(matching, dim=2)

    # has shape (num_f, num_eps, num_samples)
    return reduced