import torch


def _topk(heat: torch.Tensor, k):
    """
    Enhances the torch.topk version
    torch.topk(input, k, dim=-1) -> (values, indices)
    indices of the input tensor, values and indices
    are sorted in descending order!

    Parameters
    ----------
    heat: b x c x h x w
        model output heat map
    k: int
        find the k best model output positions

    Returns
    -------
    (topk_score, topk_inds, topk_cids, topk_ys, topk_xs)

    topk_score: b x k
        scores are values form peaks in h x w plane, over multiple
        classes c(or channels)
    topk_inds: b x k
        indices values [0, h * w), over multiple classes
    topk_cids: b x k
        [0, num. of classes), class index to which the score and inds
        belong

    -> each entry in inds can be given as (x, y) coordinate tuple
    topk_ys: b x k
    topk_xs: b x k
    """

    batch, cat, height, width = heat.size()
    #topk_scores, topk_inds = torch.topk(heat.view(batch, -1), k)
    topk_scores, topk_inds = torch.topk(heat.reshape(batch, -1), k)
    topk_cids = torch.true_divide(topk_inds, (height * width)).int()
    topk_inds = topk_inds % (height * width)
    topk_ys = torch.true_divide(topk_inds, width).int().float()
    topk_xs = (topk_inds % width).int().float()
    return topk_scores, topk_inds, topk_cids, topk_ys, topk_xs