# 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.loss.mesh_laplacian_smoothing import mesh_laplacian_smoothing
from pytorch3d.structures.meshes import Meshes


class TestLaplacianSmoothing(unittest.TestCase):
    @staticmethod
    def laplacian_smoothing_naive_uniform(meshes):
        """
        Naive implementation of laplacian smoothing with uniform weights.
        """
        verts_packed = meshes.verts_packed()  # (sum(V_n), 3)
        faces_packed = meshes.faces_packed()  # (sum(F_n), 3)
        V = verts_packed.shape[0]

        L = torch.zeros((V, V), dtype=torch.float32, device=meshes.device)

        # filling L with the face pairs should be the same as edge pairs
        for f in faces_packed:
            L[f[0], f[1]] = 1
            L[f[0], f[2]] = 1
            L[f[1], f[2]] = 1
            # symetric
            L[f[1], f[0]] = 1
            L[f[2], f[0]] = 1
            L[f[2], f[1]] = 1

        norm_w = L.sum(dim=1, keepdims=True)
        idx = norm_w > 0
        norm_w[idx] = 1.0 / norm_w[idx]

        loss = (L.mm(verts_packed) * norm_w - verts_packed).norm(dim=1)

        weights = torch.zeros(V, dtype=torch.float32, device=meshes.device)
        for v in range(V):
            weights[v] = meshes.num_verts_per_mesh()[
                meshes.verts_packed_to_mesh_idx()[v]
            ]
        weights = 1.0 / weights
        loss = loss * weights

        return loss.sum() / len(meshes)

    @staticmethod
    def laplacian_smoothing_naive_cot(meshes, method: str = "cot"):
        """
        Naive implementation of laplacian smoothing wit cotangent weights.
        """
        verts_packed = meshes.verts_packed()  # (sum(V_n), 3)
        faces_packed = meshes.faces_packed()  # (sum(F_n), 3)
        V = verts_packed.shape[0]

        L = torch.zeros((V, V), dtype=torch.float32, device=meshes.device)
        inv_areas = torch.zeros((V, 1), dtype=torch.float32, device=meshes.device)

        for f in faces_packed:
            v0 = verts_packed[f[0], :]
            v1 = verts_packed[f[1], :]
            v2 = verts_packed[f[2], :]
            A = (v1 - v2).norm()
            B = (v0 - v2).norm()
            C = (v0 - v1).norm()
            s = 0.5 * (A + B + C)

            face_area = (s * (s - A) * (s - B) * (s - C)).clamp_(min=1e-12).sqrt()
            inv_areas[f[0]] += face_area
            inv_areas[f[1]] += face_area
            inv_areas[f[2]] += face_area

            A2, B2, C2 = A * A, B * B, C * C
            cota = (B2 + C2 - A2) / face_area / 4.0
            cotb = (A2 + C2 - B2) / face_area / 4.0
            cotc = (A2 + B2 - C2) / face_area / 4.0

            L[f[1], f[2]] += cota
            L[f[2], f[0]] += cotb
            L[f[0], f[1]] += cotc
            # symetric
            L[f[2], f[1]] += cota
            L[f[0], f[2]] += cotb
            L[f[1], f[0]] += cotc

        idx = inv_areas > 0
        inv_areas[idx] = 1.0 / inv_areas[idx]

        norm_w = L.sum(dim=1, keepdims=True)
        L_sum = norm_w.clone()
        idx = norm_w > 0
        norm_w[idx] = 1.0 / norm_w[idx]

        if method == "cotcurv":
            loss = (L.mm(verts_packed) - L_sum * verts_packed) * inv_areas * 0.25
            loss = loss.norm(dim=1)
        else:
            loss = L.mm(verts_packed) * norm_w - verts_packed
            loss = loss.norm(dim=1)

        weights = torch.zeros(V, dtype=torch.float32, device=meshes.device)
        for v in range(V):
            weights[v] = meshes.num_verts_per_mesh()[
                meshes.verts_packed_to_mesh_idx()[v]
            ]
        weights = 1.0 / weights
        loss = loss * weights

        return loss.sum() / len(meshes)

    @staticmethod
    def init_meshes(num_meshes: int = 10, num_verts: int = 1000, num_faces: int = 3000):
        device = torch.device("cuda:0")
        verts_list = []
        faces_list = []
        for _ in range(num_meshes):
            verts = (
                torch.rand((num_verts, 3), dtype=torch.float32, device=device) * 2.0
                - 1.0
            )  # verts in the space of [-1, 1]
            faces = torch.stack(
                [
                    torch.randperm(num_verts, device=device)[:3]
                    for _ in range(num_faces)
                ],
                dim=0,
            )
            # avoids duplicate vertices in a face
            verts_list.append(verts)
            faces_list.append(faces)
        meshes = Meshes(verts_list, faces_list)

        return meshes

    def test_laplacian_smoothing_uniform(self):
        """
        Test Laplacian Smoothing with uniform weights.
        """
        meshes = TestLaplacianSmoothing.init_meshes(10, 100, 300)

        # feats in list
        out = mesh_laplacian_smoothing(meshes, method="uniform")
        naive_out = TestLaplacianSmoothing.laplacian_smoothing_naive_uniform(meshes)

        self.assertTrue(torch.allclose(out, naive_out))

    def test_laplacian_smoothing_cot(self):
        """
        Test Laplacian Smoothing with cot weights.
        """
        meshes = TestLaplacianSmoothing.init_meshes(10, 100, 300)

        # feats in list
        out = mesh_laplacian_smoothing(meshes, method="cot")
        naive_out = TestLaplacianSmoothing.laplacian_smoothing_naive_cot(
            meshes, method="cot"
        )

        self.assertTrue(torch.allclose(out, naive_out))

    def test_laplacian_smoothing_cotcurv(self):
        """
        Test Laplacian Smoothing with cotcurv weights.
        """
        meshes = TestLaplacianSmoothing.init_meshes(10, 100, 300)

        # feats in list
        out = mesh_laplacian_smoothing(meshes, method="cotcurv")
        naive_out = TestLaplacianSmoothing.laplacian_smoothing_naive_cot(
            meshes, method="cotcurv"
        )

        self.assertTrue(torch.allclose(out, naive_out))

    @staticmethod
    def laplacian_smoothing_with_init(
        num_meshes: int, num_verts: int, num_faces: int, 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)
        torch.cuda.synchronize()

        def smooth():
            mesh_laplacian_smoothing(meshes, method="cotcurv")
            torch.cuda.synchronize()

        return smooth