# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import unittest

import torch
from common_testing import TestCaseMixin, get_random_cuda_device
from pytorch3d.ops import packed_to_padded, padded_to_packed
from pytorch3d.structures.meshes import Meshes


class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(1)

    @staticmethod
    def init_meshes(
        num_meshes: int = 10,
        num_verts: int = 1000,
        num_faces: int = 3000,
        device: str = "cpu",
    ):
        device = torch.device(device)
        verts_list = []
        faces_list = []
        for _ in range(num_meshes):
            verts = torch.rand((num_verts, 3), dtype=torch.float32, device=device)
            faces = torch.randint(
                num_verts, size=(num_faces, 3), dtype=torch.int64, device=device
            )
            verts_list.append(verts)
            faces_list.append(faces)
        meshes = Meshes(verts_list, faces_list)

        return meshes

    @staticmethod
    def packed_to_padded_python(inputs, first_idxs, max_size, device):
        """
        PyTorch implementation of packed_to_padded function.
        """
        num_meshes = first_idxs.size(0)
        D = inputs.shape[1] if inputs.dim() == 2 else 0
        if D == 0:
            inputs_padded = torch.zeros((num_meshes, max_size), device=device)
        else:
            inputs_padded = torch.zeros((num_meshes, max_size, D), device=device)
        for m in range(num_meshes):
            s = first_idxs[m]
            if m == num_meshes - 1:
                f = inputs.shape[0]
            else:
                f = first_idxs[m + 1]
            inputs_padded[m, :f] = inputs[s:f]

        return inputs_padded

    @staticmethod
    def padded_to_packed_python(inputs, first_idxs, num_inputs, device):
        """
        PyTorch implementation of padded_to_packed function.
        """
        num_meshes = inputs.size(0)
        D = inputs.shape[2] if inputs.dim() == 3 else 0
        if D == 0:
            inputs_packed = torch.zeros((num_inputs,), device=device)
        else:
            inputs_packed = torch.zeros((num_inputs, D), device=device)
        for m in range(num_meshes):
            s = first_idxs[m]
            if m == num_meshes - 1:
                f = num_inputs
            else:
                f = first_idxs[m + 1]
            inputs_packed[s:f] = inputs[m, :f]

        return inputs_packed

    def _test_packed_to_padded_helper(self, D, device):
        """
        Check the results from packed_to_padded and PyTorch implementations
        are the same.
        """
        meshes = self.init_meshes(16, 100, 300, device=device)
        faces = meshes.faces_packed()
        mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
        max_faces = meshes.num_faces_per_mesh().max().item()

        if D == 0:
            values = torch.rand((faces.shape[0],), device=device, requires_grad=True)
        else:
            values = torch.rand((faces.shape[0], D), device=device, requires_grad=True)
        values_torch = values.detach().clone()
        values_torch.requires_grad = True
        values_padded = packed_to_padded(
            values, mesh_to_faces_packed_first_idx, max_faces
        )
        values_padded_torch = TestPackedToPadded.packed_to_padded_python(
            values_torch, mesh_to_faces_packed_first_idx, max_faces, device
        )
        # check forward
        self.assertClose(values_padded, values_padded_torch)

        # check backward
        if D == 0:
            grad_inputs = torch.rand((len(meshes), max_faces), device=device)
        else:
            grad_inputs = torch.rand((len(meshes), max_faces, D), device=device)
        values_padded.backward(grad_inputs)
        grad_outputs = values.grad
        values_padded_torch.backward(grad_inputs)
        grad_outputs_torch1 = values_torch.grad
        grad_outputs_torch2 = TestPackedToPadded.padded_to_packed_python(
            grad_inputs, mesh_to_faces_packed_first_idx, values.size(0), device=device
        )
        self.assertClose(grad_outputs, grad_outputs_torch1)
        self.assertClose(grad_outputs, grad_outputs_torch2)

    def test_packed_to_padded_flat_cpu(self):
        self._test_packed_to_padded_helper(0, "cpu")

    def test_packed_to_padded_D1_cpu(self):
        self._test_packed_to_padded_helper(1, "cpu")

    def test_packed_to_padded_D16_cpu(self):
        self._test_packed_to_padded_helper(16, "cpu")

    def test_packed_to_padded_flat_cuda(self):
        device = get_random_cuda_device()
        self._test_packed_to_padded_helper(0, device)

    def test_packed_to_padded_D1_cuda(self):
        device = get_random_cuda_device()
        self._test_packed_to_padded_helper(1, device)

    def test_packed_to_padded_D16_cuda(self):
        device = get_random_cuda_device()
        self._test_packed_to_padded_helper(16, device)

    def _test_padded_to_packed_helper(self, D, device):
        """
        Check the results from packed_to_padded and PyTorch implementations
        are the same.
        """
        meshes = self.init_meshes(16, 100, 300, device=device)
        mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
        num_faces_per_mesh = meshes.num_faces_per_mesh()
        max_faces = num_faces_per_mesh.max().item()
        if D == 0:
            values = torch.rand((len(meshes), max_faces), device=device)
        else:
            values = torch.rand((len(meshes), max_faces, D), device=device)
        for i, num in enumerate(num_faces_per_mesh):
            values[i, num:] = 0
        values.requires_grad = True
        values_torch = values.detach().clone()
        values_torch.requires_grad = True
        values_packed = padded_to_packed(
            values, mesh_to_faces_packed_first_idx, num_faces_per_mesh.sum().item()
        )
        values_packed_torch = TestPackedToPadded.padded_to_packed_python(
            values_torch,
            mesh_to_faces_packed_first_idx,
            num_faces_per_mesh.sum().item(),
            device,
        )
        # check forward
        self.assertClose(values_packed, values_packed_torch)

        # check backward
        if D == 0:
            grad_inputs = torch.rand((num_faces_per_mesh.sum().item()), device=device)
        else:
            grad_inputs = torch.rand(
                (num_faces_per_mesh.sum().item(), D), device=device
            )
        values_packed.backward(grad_inputs)
        grad_outputs = values.grad
        values_packed_torch.backward(grad_inputs)
        grad_outputs_torch1 = values_torch.grad
        grad_outputs_torch2 = TestPackedToPadded.packed_to_padded_python(
            grad_inputs, mesh_to_faces_packed_first_idx, values.size(1), device=device
        )
        self.assertClose(grad_outputs, grad_outputs_torch1)
        self.assertClose(grad_outputs, grad_outputs_torch2)

    def test_padded_to_packed_flat_cpu(self):
        self._test_padded_to_packed_helper(0, "cpu")

    def test_padded_to_packed_D1_cpu(self):
        self._test_padded_to_packed_helper(1, "cpu")

    def test_padded_to_packed_D16_cpu(self):
        self._test_padded_to_packed_helper(16, "cpu")

    def test_padded_to_packed_flat_cuda(self):
        device = get_random_cuda_device()
        self._test_padded_to_packed_helper(0, device)

    def test_padded_to_packed_D1_cuda(self):
        device = get_random_cuda_device()
        self._test_padded_to_packed_helper(1, device)

    def test_padded_to_packed_D16_cuda(self):
        device = get_random_cuda_device()
        self._test_padded_to_packed_helper(16, device)

    def test_invalid_inputs_shapes(self, device="cuda:0"):
        with self.assertRaisesRegex(ValueError, "input can only be 2-dimensional."):
            values = torch.rand((100, 50, 2), device=device)
            first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
            packed_to_padded(values, first_idxs, 100)

        with self.assertRaisesRegex(ValueError, "input can only be 3-dimensional."):
            values = torch.rand((100,), device=device)
            first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
            padded_to_packed(values, first_idxs, 20)

        with self.assertRaisesRegex(ValueError, "input can only be 3-dimensional."):
            values = torch.rand((100, 50, 2, 2), device=device)
            first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
            padded_to_packed(values, first_idxs, 20)

    @staticmethod
    def packed_to_padded_with_init(
        num_meshes: int, num_verts: int, num_faces: int, num_d: int, device: str = "cpu"
    ):
        meshes = TestPackedToPadded.init_meshes(
            num_meshes, num_verts, num_faces, device
        )
        faces = meshes.faces_packed()
        mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
        max_faces = meshes.num_faces_per_mesh().max().item()
        if num_d == 0:
            values = torch.rand((faces.shape[0],), device=meshes.device)
        else:
            values = torch.rand((faces.shape[0], num_d), device=meshes.device)
        torch.cuda.synchronize()

        def out():
            packed_to_padded(values, mesh_to_faces_packed_first_idx, max_faces)
            torch.cuda.synchronize()

        return out

    @staticmethod
    def packed_to_padded_with_init_torch(
        num_meshes: int, num_verts: int, num_faces: int, num_d: int, device: str = "cpu"
    ):
        meshes = TestPackedToPadded.init_meshes(
            num_meshes, num_verts, num_faces, device
        )
        faces = meshes.faces_packed()
        mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
        max_faces = meshes.num_faces_per_mesh().max().item()
        if num_d == 0:
            values = torch.rand((faces.shape[0],), device=meshes.device)
        else:
            values = torch.rand((faces.shape[0], num_d), device=meshes.device)
        torch.cuda.synchronize()

        def out():
            TestPackedToPadded.packed_to_padded_python(
                values, mesh_to_faces_packed_first_idx, max_faces, device
            )
            torch.cuda.synchronize()

        return out

[((18, 8, 18, 28), 'torch.manual_seed', 'torch.manual_seed', ({(18, 26, 18, 27): '(1)'}, {}), '(1)', False, 'import torch\n'), ((27, 17, 27, 37), 'torch.device', 'torch.device', ({(27, 30, 27, 36): 'device'}, {}), '(device)', False, 'import torch\n'), ((37, 17, 37, 47), 'pytorch3d.structures.meshes.Meshes', 'Meshes', ({(37, 24, 37, 34): 'verts_list', (37, 36, 37, 46): 'faces_list'}, {}), '(verts_list, faces_list)', False, 'from pytorch3d.structures.meshes import Meshes\n'), ((99, 24, 101, 9), 'pytorch3d.ops.packed_to_padded', 'packed_to_padded', ({(100, 12, 100, 18): 'values', (100, 20, 100, 50): 'mesh_to_faces_packed_first_idx', (100, 52, 100, 61): 'max_faces'}, {}), '(values, mesh_to_faces_packed_first_idx, max_faces)', False, 'from pytorch3d.ops import packed_to_padded, padded_to_packed\n'), ((133, 17, 133, 41), 'common_testing.get_random_cuda_device', 'get_random_cuda_device', ({}, {}), '()', False, 'from common_testing import TestCaseMixin, get_random_cuda_device\n'), ((137, 17, 137, 41), 'common_testing.get_random_cuda_device', 'get_random_cuda_device', ({}, {}), '()', False, 'from common_testing import TestCaseMixin, get_random_cuda_device\n'), ((141, 17, 141, 41), 'common_testing.get_random_cuda_device', 'get_random_cuda_device', ({}, {}), '()', False, 'from common_testing import TestCaseMixin, get_random_cuda_device\n'), ((201, 17, 201, 41), 'common_testing.get_random_cuda_device', 'get_random_cuda_device', ({}, {}), '()', False, 'from common_testing import TestCaseMixin, get_random_cuda_device\n'), ((205, 17, 205, 41), 'common_testing.get_random_cuda_device', 'get_random_cuda_device', ({}, {}), '()', False, 'from common_testing import TestCaseMixin, get_random_cuda_device\n'), ((209, 17, 209, 41), 'common_testing.get_random_cuda_device', 'get_random_cuda_device', ({}, {}), '()', False, 'from common_testing import TestCaseMixin, get_random_cuda_device\n'), ((242, 8, 242, 32), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((264, 8, 264, 32), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((31, 20, 31, 82), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((32, 20, 34, 13), 'torch.randint', 'torch.randint', (), '', False, 'import torch\n'), ((49, 28, 49, 78), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((51, 28, 51, 81), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((70, 28, 70, 69), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((72, 28, 72, 71), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((94, 21, 94, 85), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((96, 21, 96, 87), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((214, 21, 214, 60), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((215, 25, 215, 80), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((216, 12, 216, 53), 'pytorch3d.ops.packed_to_padded', 'packed_to_padded', ({(216, 29, 216, 35): 'values', (216, 37, 216, 47): 'first_idxs', (216, 49, 216, 52): '(100)'}, {}), '(values, first_idxs, 100)', False, 'from pytorch3d.ops import packed_to_padded, padded_to_packed\n'), ((219, 21, 219, 54), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((220, 25, 220, 80), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((221, 12, 221, 52), 'pytorch3d.ops.padded_to_packed', 'padded_to_packed', ({(221, 29, 221, 35): 'values', (221, 37, 221, 47): 'first_idxs', (221, 49, 221, 51): '(20)'}, {}), '(values, first_idxs, 20)', False, 'from pytorch3d.ops import packed_to_padded, padded_to_packed\n'), ((224, 21, 224, 63), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((225, 25, 225, 80), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((226, 12, 226, 52), 'pytorch3d.ops.padded_to_packed', 'padded_to_packed', ({(226, 29, 226, 35): 'values', (226, 37, 226, 47): 'first_idxs', (226, 49, 226, 51): '(20)'}, {}), '(values, first_idxs, 20)', False, 'from pytorch3d.ops import packed_to_padded, padded_to_packed\n'), ((239, 21, 239, 72), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((241, 21, 241, 78), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((245, 12, 245, 79), 'pytorch3d.ops.packed_to_padded', 'packed_to_padded', ({(245, 29, 245, 35): 'values', (245, 37, 245, 67): 'mesh_to_faces_packed_first_idx', (245, 69, 245, 78): 'max_faces'}, {}), '(values, mesh_to_faces_packed_first_idx, max_faces)', False, 'from pytorch3d.ops import packed_to_padded, padded_to_packed\n'), ((246, 12, 246, 36), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((261, 21, 261, 72), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((263, 21, 263, 78), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((270, 12, 270, 36), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n')]