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

import shutil

import torch
import torch.nn.functional as F
from torch.autograd import gradcheck
from torchvision.transforms import ToTensor, ToPILImage
from PIL import Image

from nemotron_ocr.inference.post_processing.quad_rectify import QuadRectify

example = 3
isotropic = True

# root_dir = '/mnt/fsx-datasets-a-1-new/mranzinger/ocr/scene-text/icdar/incidental_text/train'
root_dir = "/home/dcg-adlr-mranzinger-data.cosmos1100/ocr/scene-text/icdar/incidental_text/train"
# root_dir = '/raid/local_datasets/scene-text/icdar/focused_text/train'

image_path = "{}/images/img_{}.jpg".format(root_dir, example)
label_path = "{}/gt/gt_img_{}.txt".format(root_dir, example)

shutil.copyfile(image_path, "original.jpg")

image = Image.open(image_path)

image = ToTensor()(image)

print("image", image.shape)

quads = []

with open(label_path, "r") as fd:
    for i, line in enumerate(fd.readlines()):
        line = line.strip()
        print(line)
        # Skip the unicode character
        if i == 0:
            line = line[1:]
        line = line.split(",")
        coords = [float(t) for t in line[:8]]

        word = line[-1]

        if word != "###":
            quads.append(coords)

quads = torch.tensor(quads).reshape(-1, 4, 2)

# quads[:,:,0] /= image.shape[-1]
# quads[:,:,1] /= image.shape[-2]

# print(quads)

image = image.unsqueeze(0).repeat(quads.shape[0], 1, 1, 1)

image = image.cuda()
quads = quads.cuda()

aspect = image.shape[-1] / image.shape[-2]
qr = QuadRectify(60, 400, 0, isotropic=isotropic)
# qr = QuadRectify(60, -1000, aspect)

grid = qr(quads, *image.shape[2:])

print("grid", grid.shape)

resampled = F.grid_sample(image, grid, align_corners=False)

print(resampled.shape)

resampled = resampled.permute(1, 0, 2, 3).contiguous().reshape(3, -1, resampled.shape[-1])
resampled = resampled.cpu()

pilOutput = ToPILImage()(resampled)

pilOutput.save("rectified.jpg")

# sys.exit(0)

print("checking gradients")

quads = quads.double()
cuda_quads = quads.cuda()
quads.requires_grad_()
cuda_quads.requires_grad_()

qr = QuadRectify(9, 12, isotropic=isotropic)

print("check GPU gradients:")
# gradcheck(qr.forward, cuda_quads, eps=1e-3, atol=1e-3)
gradcheck(qr.forward, (cuda_quads, *image.shape[2:]))
print("check CPU gradients:")
gradcheck(qr.forward, (quads, *image.shape[2:]))
# #gradcheck(qr.forward, quads, eps=1e-3, atol=1e-3)
# gradcheck(qr.forward, quads)

print("checking performance...")

quads = quads.detach()

quads = torch.rand(256, 4, 2, dtype=torch.float64)
cuda_quads = quads.cuda()

qr = QuadRectify(8, 60, isotropic=isotropic)

# num_passes = 500
# for target_device in [torch.device('cuda:0'), torch.device('cpu')]:
#     for target_type in [torch.float16, torch.float32, torch.float64]:
#         test_quads = quads.to(target_type).to(target_device).clone()

#         test_quads.requires_grad_()

#         print('type:', target_type, 'device:', target_device)

#         try:
#             with torch.autograd.profiler.profile(use_cuda=test_quads.is_cuda) as prof:
#                 for _ in range(num_passes):
#                     g = qr(test_quads)
#                     l = (1 - g.mean())
#                     l.backward()
#             print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

#             # print('\tforward...')
#             # torch.cuda.synchronize()
#             # start_time = time.time()
#             # for i in range(num_passes):
#             #     qr(test_quads).sum()
#             #     torch.cuda.synchronize()
#             # end_time = time.time()
#             # fp_sec_per_call = (end_time - start_time) / num_passes
#             # print('\t\tTook', fp_sec_per_call, "sec/call")

#             # print('\tbackward...')
#             # torch.cuda.synchronize()
#             # start_time = time.time()
#             # for i in range(num_passes):
#             #     g = qr(test_quads)
#             #     g.sum().backward()
#             #     torch.cuda.synchronize()
#             # end_time = time.time()
#             # bp_sec_per_call = (end_time - start_time) / num_passes - fp_sec_per_call
#             # print('\t\tTook', bp_sec_per_call, "sec/call")
#         except Exception as e:
#             print('\t\t', e)

# # for target_type in [torch.float16, torch.float32]:
# #     print('DType:', target_type)
# #     h_cuda_quads = cuda_quads.clone().to(target_type)
# #     h_cuda_quads.requires_grad_()

# #     v = qr(h_cuda_quads).sum()
# #     v.backward()

# #     print(h_cuda_quads.grad)