import torch


def bagnet_predict(model, images, k=1, clip=None, a=1e-2, b=-0.78, return_class=True):
    """ Make top-K prediction on IMAGES by MODEL
    Inputs:
    - model: pytorch model. model for prediction
    - images: pytorch tensor. images to be predicted on
    - k: number of classes to return for each image (top-k most possible ones)
    - clip (clipping): clipping function
    - a, b (double): clipping parameters
    - return_class: If True, then return classes as prediction; otherwise, return logits and probability of prediction classes
    Return:
    - indices.cpu().numpy(): numpy array at CPU. prediction K classes
    - l.cpu().numpy(): numpy array at CPU. top-K prediction logits
    - p.cpu().numpy(): numpy array at CPU. top-K prediction probability
    """
    with torch.no_grad():
        logits = model(images)
        if clip:
            assert not model.avg_pool, 'Bagnet should apply clipping before taking average.'
            logits = clip(logits, a, b)
            logits = torch.mean(logits, dim=(1, 2))
        p = torch.nn.Softmax(dim=1)(logits)
        p, indices = torch.topk(p, k, dim=1)
        l, _ = torch.topk(logits, k, dim=1)
        if return_class:
            return indices.cpu().numpy()
        else:
            return l.cpu().numpy(), p.cpu().numpy()