# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import unittest
import torch

from pytorch3d.ops import cubify


class TestCubify(unittest.TestCase):
    def test_allempty(self):
        N, V = 32, 14
        device = torch.device("cuda:0")
        voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)
        meshes = cubify(voxels, 0.5, 0)
        self.assertTrue(meshes.isempty)

    def test_cubify(self):
        N, V = 4, 2
        device = torch.device("cuda:0")
        voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)

        # 1st example: (top left corner, znear) is on
        voxels[0, 0, 0, 0] = 1.0
        # 2nd example: all are on
        voxels[1] = 1.0
        # 3rd example: empty
        # 4th example
        voxels[3, :, :, 1] = 1.0
        voxels[3, 1, 1, 0] = 1.0

        # compute cubify
        meshes = cubify(voxels, 0.5, 0)

        # 1st-check
        verts, faces = meshes.get_mesh_verts_faces(0)
        self.assertTrue(
            torch.allclose(faces.max(), torch.tensor([verts.size(0) - 1]))
        )
        self.assertTrue(
            torch.allclose(
                verts,
                torch.tensor(
                    [
                        [-1.0, -1.0, -1.0],
                        [-1.0, -1.0, 1.0],
                        [1.0, -1.0, -1.0],
                        [1.0, -1.0, 1.0],
                        [-1.0, 1.0, -1.0],
                        [-1.0, 1.0, 1.0],
                        [1.0, 1.0, -1.0],
                        [1.0, 1.0, 1.0],
                    ],
                    dtype=torch.float32,
                    device=device,
                ),
            )
        )
        self.assertTrue(
            torch.allclose(
                faces,
                torch.tensor(
                    [
                        [0, 1, 4],
                        [1, 5, 4],
                        [4, 5, 6],
                        [5, 7, 6],
                        [0, 4, 6],
                        [0, 6, 2],
                        [0, 3, 1],
                        [0, 2, 3],
                        [6, 7, 3],
                        [6, 3, 2],
                        [1, 7, 5],
                        [1, 3, 7],
                    ],
                    dtype=torch.int64,
                    device=device,
                ),
            )
        )
        # 2nd-check
        verts, faces = meshes.get_mesh_verts_faces(1)
        self.assertTrue(
            torch.allclose(faces.max(), torch.tensor([verts.size(0) - 1]))
        )
        self.assertTrue(
            torch.allclose(
                verts,
                torch.tensor(
                    [
                        [-1.0, -1.0, -1.0],
                        [-1.0, -1.0, 1.0],
                        [-1.0, -1.0, 3.0],
                        [1.0, -1.0, -1.0],
                        [1.0, -1.0, 1.0],
                        [1.0, -1.0, 3.0],
                        [3.0, -1.0, -1.0],
                        [3.0, -1.0, 1.0],
                        [3.0, -1.0, 3.0],
                        [-1.0, 1.0, -1.0],
                        [-1.0, 1.0, 1.0],
                        [-1.0, 1.0, 3.0],
                        [1.0, 1.0, -1.0],
                        [1.0, 1.0, 3.0],
                        [3.0, 1.0, -1.0],
                        [3.0, 1.0, 1.0],
                        [3.0, 1.0, 3.0],
                        [-1.0, 3.0, -1.0],
                        [-1.0, 3.0, 1.0],
                        [-1.0, 3.0, 3.0],
                        [1.0, 3.0, -1.0],
                        [1.0, 3.0, 1.0],
                        [1.0, 3.0, 3.0],
                        [3.0, 3.0, -1.0],
                        [3.0, 3.0, 1.0],
                        [3.0, 3.0, 3.0],
                    ],
                    dtype=torch.float32,
                    device=device,
                ),
            )
        )
        self.assertTrue(
            torch.allclose(
                faces,
                torch.tensor(
                    [
                        [0, 1, 9],
                        [1, 10, 9],
                        [0, 9, 12],
                        [0, 12, 3],
                        [0, 4, 1],
                        [0, 3, 4],
                        [1, 2, 10],
                        [2, 11, 10],
                        [1, 5, 2],
                        [1, 4, 5],
                        [2, 13, 11],
                        [2, 5, 13],
                        [3, 12, 14],
                        [3, 14, 6],
                        [3, 7, 4],
                        [3, 6, 7],
                        [14, 15, 7],
                        [14, 7, 6],
                        [4, 8, 5],
                        [4, 7, 8],
                        [15, 16, 8],
                        [15, 8, 7],
                        [5, 16, 13],
                        [5, 8, 16],
                        [9, 10, 17],
                        [10, 18, 17],
                        [17, 18, 20],
                        [18, 21, 20],
                        [9, 17, 20],
                        [9, 20, 12],
                        [10, 11, 18],
                        [11, 19, 18],
                        [18, 19, 21],
                        [19, 22, 21],
                        [11, 22, 19],
                        [11, 13, 22],
                        [20, 21, 23],
                        [21, 24, 23],
                        [12, 20, 23],
                        [12, 23, 14],
                        [23, 24, 15],
                        [23, 15, 14],
                        [21, 22, 24],
                        [22, 25, 24],
                        [24, 25, 16],
                        [24, 16, 15],
                        [13, 25, 22],
                        [13, 16, 25],
                    ],
                    dtype=torch.int64,
                    device=device,
                ),
            )
        )

        # 3rd-check
        verts, faces = meshes.get_mesh_verts_faces(2)
        self.assertTrue(verts.size(0) == 0)
        self.assertTrue(faces.size(0) == 0)

        # 4th-check
        verts, faces = meshes.get_mesh_verts_faces(3)
        self.assertTrue(
            torch.allclose(
                verts,
                torch.tensor(
                    [
                        [1.0, -1.0, -1.0],
                        [1.0, -1.0, 1.0],
                        [1.0, -1.0, 3.0],
                        [3.0, -1.0, -1.0],
                        [3.0, -1.0, 1.0],
                        [3.0, -1.0, 3.0],
                        [-1.0, 1.0, 1.0],
                        [-1.0, 1.0, 3.0],
                        [1.0, 1.0, -1.0],
                        [1.0, 1.0, 1.0],
                        [1.0, 1.0, 3.0],
                        [3.0, 1.0, -1.0],
                        [3.0, 1.0, 1.0],
                        [3.0, 1.0, 3.0],
                        [-1.0, 3.0, 1.0],
                        [-1.0, 3.0, 3.0],
                        [1.0, 3.0, -1.0],
                        [1.0, 3.0, 1.0],
                        [1.0, 3.0, 3.0],
                        [3.0, 3.0, -1.0],
                        [3.0, 3.0, 1.0],
                        [3.0, 3.0, 3.0],
                    ],
                    dtype=torch.float32,
                    device=device,
                ),
            )
        )
        self.assertTrue(
            torch.allclose(
                faces,
                torch.tensor(
                    [
                        [0, 1, 8],
                        [1, 9, 8],
                        [0, 8, 11],
                        [0, 11, 3],
                        [0, 4, 1],
                        [0, 3, 4],
                        [11, 12, 4],
                        [11, 4, 3],
                        [1, 2, 9],
                        [2, 10, 9],
                        [1, 5, 2],
                        [1, 4, 5],
                        [12, 13, 5],
                        [12, 5, 4],
                        [2, 13, 10],
                        [2, 5, 13],
                        [6, 7, 14],
                        [7, 15, 14],
                        [14, 15, 17],
                        [15, 18, 17],
                        [6, 14, 17],
                        [6, 17, 9],
                        [6, 10, 7],
                        [6, 9, 10],
                        [7, 18, 15],
                        [7, 10, 18],
                        [8, 9, 16],
                        [9, 17, 16],
                        [16, 17, 19],
                        [17, 20, 19],
                        [8, 16, 19],
                        [8, 19, 11],
                        [19, 20, 12],
                        [19, 12, 11],
                        [17, 18, 20],
                        [18, 21, 20],
                        [20, 21, 13],
                        [20, 13, 12],
                        [10, 21, 18],
                        [10, 13, 21],
                    ],
                    dtype=torch.int64,
                    device=device,
                ),
            )
        )

    @staticmethod
    def cubify_with_init(batch_size: int, V: int):
        device = torch.device("cuda:0")
        voxels = torch.rand(
            (batch_size, V, V, V), dtype=torch.float32, device=device
        )
        torch.cuda.synchronize()

        def convert():
            cubify(voxels, 0.5)
            torch.cuda.synchronize()

        return convert