import torch


def pick_targeted_classes(logits, k=5, case='avg'):
    """
    Input:
    - logits (pytorch tensor): clipped logits (clipped patch logits after avg). Shape (N, 1000)
    - k (int): top-k prediction
    - case (string): ['best', 'avg', 'worst']
            'best': the class ranked at (k+1)-th
            'avg': randomly sample one class outside of the top-k predictions
            'worst': the class ranked the last
    """
    assert case in ['best', 'avg', 'worst'], 'case must be "best", "avg", or "worst"'
    if case == 'best':
        _, topk = torch.topk(logits, k=k+1, dim=1) # shape (N, k+1) 
        return topk[:, -1]
    if case == 'avg':
        _, topk = torch.topk(logits, k=1000, dim=1) # shape (N, 1000) 
        indices = torch.randint(low=5, high=999, size=(1,))
        return torch.index_select(topk, 1, indices.cuda()).reshape((-1,))
    else: # if case == 'worst'
        _, topk = torch.topk(logits, k=1000, dim=1) # shape (N, 1000) 
        return topk[:, -1]