def test_box_nms_op():
    def test_box_nms_forward(data, expected, thresh=0.5, valid=0, topk=-1, coord=2, score=1, cid=0, bid=-1,
                             force=False, in_format='corner', out_format='corner'):
        for dtype in ['float16', 'float32', 'float64']:
            data = mx.nd.array(data, dtype=dtype)
            out = mx.contrib.nd.box_nms(data, overlap_thresh=thresh, valid_thresh=valid, topk=topk,
                                        coord_start=coord, score_index=score, id_index=cid, background_id=bid,
                                        force_suppress=force, in_format=in_format, out_format=out_format)
            assert_almost_equal(out.asnumpy(), expected.astype(dtype), rtol=1e-3, atol=1e-3)

    def test_box_nms_backward(data, grad, expected, thresh=0.5, valid=0, topk=-1, coord=2, score=1,
                              cid=0, bid=-1, force=False, in_format='corner', out_format='corner'):
        in_var = mx.sym.Variable('data')
        arr_data = mx.nd.array(data)
        arr_grad = mx.nd.empty(arr_data.shape)
        op = mx.contrib.sym.box_nms(in_var, overlap_thresh=thresh, valid_thresh=valid, topk=topk,
                                    coord_start=coord, score_index=score, id_index=cid, background_id=bid,
                                    force_suppress=force, in_format=in_format, out_format=out_format)
        exe = op.bind(ctx=default_context(), args=[arr_data], args_grad=[arr_grad])
        exe.forward(is_train=True)
        exe.backward(mx.nd.array(grad))
        assert_almost_equal(arr_grad.asnumpy(), expected)

    def corner_to_center(data):
        out = np.reshape(data, (-1, 6)).copy()
        out[:, 2] = (data[:, 2] + data[:, 4]) / 2.0
        out[:, 3] = (data[:, 3] + data[:, 5]) / 2.0
        out[:, 4] = data[:, 4] - data[:, 2]
        out[:, 5] = data[:, 5] - data[:, 3]
        invalid = np.where(data[:, 0] < 0)[0]
        out[invalid, :] = -1
        return out

    def center_to_corner(data):
        data = np.reshape(data, (-1, 6)).copy()
        out[:, 2] = data[:, 2] - data[:, 4] / 2.0
        out[:, 3] = data[:, 3] - data[:, 5] / 2.0
        out[:, 4] = data[:, 2] + data[:, 4] / 2.0
        out[:, 5] = data[:, 3] + data[:, 5] / 2.0
        invalid = np.where(data[:, 0] < 0)[0]
        out[invalid, :] = -1
        return out

    def swap_position(data, expected, coord=2, score=1, cid=0, new_col=0):
        data = np.reshape(data, (-1, 6))
        expected = np.reshape(expected, (-1, 6))
        new_coord = random.randint(0, 6 + new_col - 4)
        others = list(range(new_coord)) + list(range(new_coord + 4, 6 + new_col))
        random.shuffle(others)
        new_score = others[0]
        new_cid = others[1]
        new_data = np.full((data.shape[0], data.shape[1] + new_col), -1.0)
        new_expected = np.full((expected.shape[0], expected.shape[1] + new_col), -1.0)
        new_data[:, new_coord:new_coord+4] = data[:, coord:coord+4]
        new_data[:, new_score] = data[:, score]
        new_data[:, new_cid] = data[:, cid]
        new_expected[:, new_coord:new_coord+4] = expected[:, coord:coord+4]
        new_expected[:, new_score] = expected[:, score]
        new_expected[:, new_cid] = expected[:, cid]
        return new_data, new_expected, new_coord, new_score, new_cid

    # manually set up test cases
    boxes = [[0, 0.5, 0.1, 0.1, 0.2, 0.2], [1, 0.4, 0.1, 0.1, 0.2, 0.2],
             [0, 0.3, 0.1, 0.1, 0.14, 0.14], [2, 0.6, 0.5, 0.5, 0.7, 0.8]]

    # case1
    force = True
    thresh = 0.5
    expected = [[2, 0.6, 0.5, 0.5, 0.7, 0.8], [0, 0.5, 0.1, 0.1, 0.2, 0.2],
                [0, 0.3, 0.1, 0.1, 0.14, 0.14], [-1, -1, -1, -1, -1, -1]]
    grad = np.random.rand(4, 6)
    expected_in_grad = grad[(1, 3, 2, 0), :]
    expected_in_grad[1, :] = 0
    test_box_nms_forward(np.array(boxes), np.array(expected), force=force, thresh=thresh)
    test_box_nms_backward(np.array(boxes), grad, expected_in_grad, force=force, thresh=thresh)

    # case2: multi batch
    boxes2 = [boxes] * 3
    expected2 = [expected] * 3
    grad2 = np.array([grad.tolist()] * 3)
    expected_in_grad2 = np.array([expected_in_grad.tolist()] * 3)
    test_box_nms_forward(np.array(boxes2), np.array(expected2), force=force, thresh=thresh)
    test_box_nms_backward(np.array(boxes2), grad2, expected_in_grad2, force=force, thresh=thresh)
    # another new dim
    boxes2 = [boxes2] * 2
    expected2 = [expected2] * 2
    grad2 = np.array([grad2.tolist()] * 2)
    expected_in_grad2 = np.array([expected_in_grad2.tolist()] * 2)
    test_box_nms_forward(np.array(boxes2), np.array(expected2), force=force, thresh=thresh)
    test_box_nms_backward(np.array(boxes2), grad2, expected_in_grad2, force=force, thresh=thresh)

    # case3: thresh
    thresh = 0.1
    boxes3 = boxes
    expected3 = [[2, 0.6, 0.5, 0.5, 0.7, 0.8], [0, 0.5, 0.1, 0.1, 0.2, 0.2],
                [-1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1]]
    grad3 = np.random.rand(4, 6)
    expected_in_grad3 = grad3[(1, 3, 2, 0), :]
    expected_in_grad3[(1, 2), :] = 0
    test_box_nms_forward(np.array(boxes3), np.array(expected3), force=force, thresh=thresh)
    test_box_nms_backward(np.array(boxes3), grad3, expected_in_grad3, force=force, thresh=thresh)

    # case4: non-force
    boxes4 = boxes
    force = False
    expected4 = [[2, 0.6, 0.5, 0.5, 0.7, 0.8], [0, 0.5, 0.1, 0.1, 0.2, 0.2],
                [1, 0.4, 0.1, 0.1, 0.2, 0.2], [-1, -1, -1, -1, -1, -1]]
    grad4 = np.random.rand(4, 6)
    expected_in_grad4 = grad4[(1, 2, 3, 0), :]
    expected_in_grad4[2, :] = 0
    test_box_nms_forward(np.array(boxes4), np.array(expected4), force=force, thresh=thresh)
    test_box_nms_backward(np.array(boxes4), grad4, expected_in_grad4, force=force, thresh=thresh)

    # case5: different coding
    boxes5 = corner_to_center(np.array(boxes4))
    test_box_nms_forward(np.array(boxes5), np.array(expected4), force=force, thresh=thresh,
        in_format='center')
    expected5 = corner_to_center(np.array(expected4))
    test_box_nms_forward(np.array(boxes4), np.array(expected5), force=force, thresh=thresh,
        out_format='center')
    test_box_nms_forward(np.array(boxes5), np.array(expected5), force=force, thresh=thresh,
        in_format='center', out_format='center')

    # case6: different position
    boxes6, expected6, new_coord, new_score, new_id = swap_position(np.array(boxes4),
        np.array(expected4), new_col=2)
    test_box_nms_forward(np.array(boxes6), np.array(expected6), force=force, thresh=thresh,
        coord=new_coord, score=new_score, cid=new_id)

    # case7: no id, should be same with force=True
    force = False
    thresh = 0.5
    test_box_nms_forward(np.array(boxes), np.array(expected), force=force, thresh=thresh, cid=-1)

    # case8: multi-batch thresh + topk
    boxes8 = [[[1, 1, 0, 0, 10, 10], [1, 0.4, 0, 0, 10, 10], [1, 0.3, 0, 0, 10, 10]],
              [[2, 1, 0, 0, 10, 10], [2, 0.4, 0, 0, 10, 10], [2, 0.3, 0, 0, 10, 10]],
              [[3, 1, 0, 0, 10, 10], [3, 0.4, 0, 0, 10, 10], [3, 0.3, 0, 0, 10, 10]]]
    expected8 = [[[1, 1, 0, 0, 10, 10], [-1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1]],
                 [[2, 1, 0, 0, 10, 10], [-1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1]],
                 [[3, 1, 0, 0, 10, 10], [-1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1]]]
    grad8 = np.random.rand(3, 3, 6)
    expected_in_grad8 = np.zeros((3, 3, 6))
    expected_in_grad8[(0, 1, 2), (0, 0, 0), :] = grad8[(0, 1, 2), (0, 0, 0), :]
    force = False
    thresh = 0.5
    valid = 0.5
    topk = 2
    test_box_nms_forward(np.array(boxes8), np.array(expected8), force=force, thresh=thresh, valid=valid, topk=topk)
    test_box_nms_backward(np.array(boxes8), grad8, expected_in_grad8, force=force, thresh=thresh, valid=valid, topk=topk)

    # case9: background id filter out
    # default background id -1
    boxes9 = [[0, 0.5, 0.1, 0.1, 0.2, 0.2], [0, 0.4, 0.1, 0.1, 0.2, 0.2],
              [1, 0.3, 0.1, 0.1, 0.14, 0.14], [-1, 0.6, 0.5, 0.5, 0.7, 0.8]]
    expected9 = [[0, 0.5, 0.1, 0.1, 0.2, 0.2], [1, 0.3, 0.1, 0.1, 0.14, 0.14],
                 [-1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1]]
    force = True
    thresh = 0.5
    grad9 = np.random.rand(4, 6)
    expected_in_grad9 = grad9[(0, 2, 1, 3), :]
    expected_in_grad9[(1, 3), :] = 0
    test_box_nms_forward(np.array(boxes9), np.array(expected9), force=force, thresh=thresh)
    test_box_nms_backward(np.array(boxes9), grad9, expected_in_grad9, force=force, thresh=thresh)
    # set background id
    background_id = 0
    expected9 = [[-1, 0.6, 0.5, 0.5, 0.7, 0.8], [1, 0.3, 0.1, 0.1, 0.14, 0.14],
                 [-1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1]]
    grad9 = np.random.rand(4, 6)
    expected_in_grad9 = grad9[(2, 3, 1, 0), :]
    expected_in_grad9[(0, 1), :] = 0
    test_box_nms_forward(np.array(boxes9), np.array(expected9), force=force, thresh=thresh, bid=background_id)
    test_box_nms_backward(np.array(boxes9), grad9, expected_in_grad9, force=force, thresh=thresh, bid=background_id)