import torch


def get_targeted_classes(model, images, clip, a, b, k=5, case='avg'):
    """ 
    Input:
    - model (pytorch model): bagnet model without avgerage pooling
    - images (pytorch tensor): a batch of preprocessed images, same context as model's images 
    - k (int): get targeted classes based on 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
    Output:
    - targeted class (pytorch tensor): targeted class for targeted attack, same context as model's.
    """
    assert case in ['best', 'avg', 'worst'], 'case must be "best", "avg", or "worst"'
    logits = model(images)
    if clip:
        logits = clip(logits, a, b)
    logits = torch.mean(logits, dim=(1, 2)) # shape (N, 1000)
    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]