# Copyright (c) Meta Platforms, Inc. and 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 pytorch3d.ops import cubify
from pytorch3d.renderer.mesh.textures import TexturesAtlas

from .common_testing import TestCaseMixin


class TestCubify(TestCaseMixin, 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)
        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)

        # 1st-check
        verts, faces = meshes.get_mesh_verts_faces(0)
        self.assertClose(faces.max().cpu(), torch.tensor(verts.size(0) - 1))
        self.assertClose(
            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.assertClose(
            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.assertClose(faces.max().cpu(), torch.tensor(verts.size(0) - 1))
        self.assertClose(
            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.assertClose(
            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.assertClose(
            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.assertClose(
            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,
            ),
        )

    def test_align(self):
        N, V = 1, 2
        device = torch.device("cuda:0")
        voxels = torch.ones((N, V, V, V), dtype=torch.float32, device=device)

        # topleft align
        mesh = cubify(voxels, 0.5)
        verts, faces = mesh.get_mesh_verts_faces(0)
        self.assertClose(verts.min(), torch.tensor(-1.0, device=device))
        self.assertClose(verts.max(), torch.tensor(3.0, device=device))

        # corner align
        mesh = cubify(voxels, 0.5, align="corner")
        verts, faces = mesh.get_mesh_verts_faces(0)
        self.assertClose(verts.min(), torch.tensor(-1.0, device=device))
        self.assertClose(verts.max(), torch.tensor(1.0, device=device))

        # center align
        mesh = cubify(voxels, 0.5, align="center")
        verts, faces = mesh.get_mesh_verts_faces(0)
        self.assertClose(verts.min(), torch.tensor(-2.0, device=device))
        self.assertClose(verts.max(), torch.tensor(2.0, device=device))

        # invalid align
        with self.assertRaisesRegex(ValueError, "Align mode must be one of"):
            cubify(voxels, 0.5, align="")

        # invalid align
        with self.assertRaisesRegex(ValueError, "Align mode must be one of"):
            cubify(voxels, 0.5, align="topright")

        # inside occupancy, similar to GH#185 use case
        N, V = 1, 4
        voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)
        voxels[0, : V // 2, : V // 2, : V // 2] = 1.0
        mesh = cubify(voxels, 0.5, align="corner")
        verts, faces = mesh.get_mesh_verts_faces(0)
        self.assertClose(verts.min(), torch.tensor(-1.0, device=device))
        self.assertClose(verts.max(), torch.tensor(0.0, 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

    def test_cubify_with_feats(self):
        N, V = 3, 2
        device = torch.device("cuda:0")
        voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)
        feats = torch.zeros((N, 3, V, V, V), dtype=torch.float32, device=device)
        # fill the feats with red color
        feats[:, 0, :, :, :] = 255

        # 1st example: (top left corner, znear) is on
        voxels[0, 0, 0, 0] = 1.0
        # the color is set to green
        feats[0, :, 0, 0, 0] = torch.Tensor([0, 255, 0])
        # 2nd example: all are on
        voxels[1] = 1.0

        # 3rd example
        voxels[2, :, :, 1] = 1.0
        voxels[2, 1, 1, 0] = 1.0
        # the color is set to yellow and blue respectively
        feats[2, 1, :, :, 1] = 255
        feats[2, :, 1, 1, 0] = torch.Tensor([0, 0, 255])
        meshes = cubify(voxels, 0.5, feats=feats, align="center")
        textures = meshes.textures
        self.assertTrue(textures is not None)
        self.assertTrue(isinstance(textures, TexturesAtlas))
        faces_textures = textures.faces_verts_textures_packed()
        red = faces_textures.new_tensor([255.0, 0.0, 0.0])
        green = faces_textures.new_tensor([0.0, 255.0, 0.0])
        blue = faces_textures.new_tensor([0.0, 0.0, 255.0])
        yellow = faces_textures.new_tensor([255.0, 255.0, 0.0])

        self.assertEqual(faces_textures.shape, (100, 3, 3))
        faces_textures_ = faces_textures.flatten(end_dim=1)
        self.assertClose(faces_textures_[:36], green.expand(36, -1))
        self.assertClose(faces_textures_[36:180], red.expand(144, -1))
        self.assertClose(faces_textures_[180:228], yellow.expand(48, -1))
        self.assertClose(faces_textures_[228:258], blue.expand(30, -1))
        self.assertClose(faces_textures_[258:300], yellow.expand(42, -1))