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

import math

import torch
from torch.autograd import gradcheck


from nemotron_ocr.inference.post_processing.rrect_to_quads import RRectToQuads


def get_rrects(b, h, w):
    rrects = torch.rand(b, h, w, 5, dtype=torch.float64)
    rrects[:, :, :, :4] *= 10
    rrects[:, :, :, 4] *= 2 * math.pi
    return rrects


rrects = get_rrects(2, 5, 5)

cell_size = 4
r2q = RRectToQuads(cell_size)

quads = r2q(rrects)

rrects.requires_grad_()

print("check CPU gradients")
gradcheck(r2q.forward, rrects)

rrects = get_rrects(4, 10, 10)
rrects.requires_grad_()

print("check GPU gradients")
gradcheck(r2q.forward, rrects)