# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import torch

from nemotron_ocr_cpp import quad_non_maximal_suppression

quads = [
    [[0, 2], [2, 2], [2, 0], [0, 0]],
    [[1, 3], [3, 3], [3, 1], [1, 1]],
    [[6, 6], [10, 6], [10, 5], [6, 5]],
    [[6, 5.5], [8, 6], [10, 5.5], [8, 5]],
    [[7, 6], [9, 6], [9, 5], [7, 5]],
]

quads = torch.tensor([[quads]], dtype=torch.float32).cuda()

probs = [0.7, 0.8, 0.55, 0.9, 0.85]

probs = torch.tensor([[probs]], dtype=torch.float32).cuda()

print("in_quads")
print(quads)
print("in_probs")
print(probs)

out_quads, out_probs, out_region_counts = quad_non_maximal_suppression(
    quads, probs, prob_threshold=0.5, iou_threshold=0.1, kernel_height=2, kernel_width=2
)[:3]

print("out_quads")
print(out_quads)
print("out_probs")
print(out_probs)
print("out_region_counts")
print(out_region_counts)