import torch

from mmaction.ops.nms import nms_wrapper


def multiclass_nms(multi_bboxes, multi_scores, score_thr, nms_cfg, max_num=-1):
    """NMS for multi-class bboxes.

    Args:
        multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
        multi_scores (Tensor): shape (n, #class)
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_thr (float): NMS IoU threshold
        max_num (int): if there are more than max_num bboxes after NMS,
            only top max_num will be kept.

    Returns:
        tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). labels
            are 0-based.
    """
    num_classes = multi_scores.shape[1]
    bboxes, labels = [], []
    nms_cfg_ = nms_cfg.copy()
    nms_type = nms_cfg_.pop('type', 'nms')
    nms_op = getattr(nms_wrapper, nms_type)
    for i in range(1, num_classes):
        cls_inds = multi_scores[:, i] > score_thr
        if not cls_inds.any():
            continue
        # get bboxes and scores of this class
        if multi_bboxes.shape[1] == 4:
            _bboxes = multi_bboxes[cls_inds, :]
        else:
            _bboxes = multi_bboxes[cls_inds, i * 4:(i + 1) * 4]
        _scores = multi_scores[cls_inds, i]
        cls_dets = torch.cat([_bboxes, _scores[:, None]], dim=1)
        cls_dets, _ = nms_op(cls_dets, **nms_cfg_)
        cls_labels = multi_bboxes.new_full(
            (cls_dets.shape[0], ), i - 1, dtype=torch.long)
        bboxes.append(cls_dets)
        labels.append(cls_labels)
    if bboxes:
        bboxes = torch.cat(bboxes)
        labels = torch.cat(labels)
        if bboxes.shape[0] > max_num:
            _, inds = bboxes[:, -1].sort(descending=True)
            inds = inds[:max_num]
            bboxes = bboxes[inds]
            labels = labels[inds]
    else:
        bboxes = multi_bboxes.new_zeros((0, 5))
        labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)

    return bboxes, labels


def singleclass_nms(multi_bboxes, multi_scores, score_thr, nms_cfg,
                    max_num=-1):
    """NMS for single-class bboxes.

    Args:
        multi_bboxes (Tensor): shape (n, 4)
        multi_scores (Tensor): shape (n, #class)
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_thr (float): NMS IoU threshold
        max_num (int): if there are more than max_num bboxes after NMS,
            only top max_num will be kept.

    Returns:
        tuple: (bboxes, scores), tensors of shape (k, 5) and (k, #class).
            labels are 0-based.
    """
    bboxes, scores = [], []
    nms_cfg_ = nms_cfg.copy()
    nms_type = nms_cfg_.pop('type', 'nms')
    nms_op = getattr(nms_wrapper, nms_type)

    cls_inds = multi_scores[:, 0] > score_thr
    if not cls_inds.any():
        bboxes = multi_bboxes.new_zeros((0, 5))
        scores = multi_bboxes.new_zeros((0, multi_scores.size(1)))
        return bboxes, scores
    # get bboxes and scores of this class
    _bboxes = multi_bboxes[cls_inds, :]
    _scores = multi_scores[cls_inds, :]
    cls_dets = torch.cat([_bboxes, _scores[:, 0:1]], dim=1)
    cls_dets, nms_keep = nms_op(cls_dets, **nms_cfg_)
    cls_scores = _scores[nms_keep, :]
    bboxes.append(cls_dets)
    scores.append(cls_scores)

    if bboxes:
        bboxes = torch.cat(bboxes)
        scores = torch.cat(scores)
        if bboxes.shape[0] > max_num:
            _, inds = bboxes[:, -1].sort(descending=True)
            inds = inds[:max_num]
            bboxes = bboxes[inds]
            scores = scores[inds]
    else:
        bboxes = multi_bboxes.new_zeros((0, 5))
        scores = multi_bboxes.new_zeros((0, multi_scores.size(1)))

    return bboxes, scores
