# 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.implicitron.tools.point_cloud_utils import get_rgbd_point_cloud

from pytorch3d.renderer.cameras import PerspectiveCameras
from tests.common_testing import TestCaseMixin


class TestPointCloudUtils(TestCaseMixin, unittest.TestCase):
    def setUp(self):
        torch.manual_seed(42)

    def test_unproject(self):
        H, W = 50, 100

        # Random RGBD image with depth 3
        # (depth 0 = at the camera)
        # and purple in the upper right corner

        image = torch.rand(4, H, W)
        depth = 3
        image[3] = depth
        image[1, H // 2 :, W // 2 :] *= 0.4

        # two ways to define the same camera:
        # at the origin facing the positive z axis
        ndc_camera = PerspectiveCameras(focal_length=1.0)
        screen_camera = PerspectiveCameras(
            focal_length=H // 2,
            in_ndc=False,
            image_size=((H, W),),
            principal_point=((W / 2, H / 2),),
        )

        for camera in (ndc_camera, screen_camera):
            # 1. z-depth
            cloud = get_rgbd_point_cloud(
                camera,
                image_rgb=image[:3][None],
                depth_map=image[3:][None],
                euclidean=False,
            )
            [points] = cloud.points_list()
            self.assertConstant(points[:, 2], depth)  # constant depth
            extremes = depth * torch.tensor([W / H - 1 / H, 1 - 1 / H])
            self.assertClose(points[:, :2].min(0).values, -extremes)
            self.assertClose(points[:, :2].max(0).values, extremes)

            # 2. euclidean
            cloud = get_rgbd_point_cloud(
                camera,
                image_rgb=image[:3][None],
                depth_map=image[3:][None],
                euclidean=True,
            )
            [points] = cloud.points_list()
            self.assertConstant(torch.norm(points, dim=1), depth, atol=1e-5)

            # 3. four channels
            get_rgbd_point_cloud(
                camera,
                image_rgb=image[None],
                depth_map=image[3:][None],
                euclidean=True,
            )