# 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 numpy as np
import torch
from pytorch3d.ops import corresponding_cameras_alignment
from pytorch3d.renderer.cameras import (
    OpenGLOrthographicCameras,
    OpenGLPerspectiveCameras,
    SfMOrthographicCameras,
    SfMPerspectiveCameras,
)
from pytorch3d.transforms.rotation_conversions import random_rotations
from pytorch3d.transforms.so3 import so3_exp_map, so3_relative_angle

from .common_testing import TestCaseMixin
from .test_cameras import init_random_cameras


class TestCamerasAlignment(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(42)
        np.random.seed(42)

    def test_corresponding_cameras_alignment(self):
        """
        Checks the corresponding_cameras_alignment function.
        """
        device = torch.device("cuda:0")

        # try few different random setups
        for _ in range(3):
            for estimate_scale in (True, False):
                # init true alignment transform
                R_align_gt = random_rotations(1, device=device)[0]
                T_align_gt = torch.randn(3, dtype=torch.float32, device=device)

                # init true scale
                if estimate_scale:
                    s_align_gt = torch.randn(
                        1, dtype=torch.float32, device=device
                    ).exp()
                else:
                    s_align_gt = torch.tensor(1.0, dtype=torch.float32, device=device)

                for cam_type in (
                    SfMOrthographicCameras,
                    OpenGLPerspectiveCameras,
                    OpenGLOrthographicCameras,
                    SfMPerspectiveCameras,
                ):
                    # try well-determined and underdetermined cases
                    for batch_size in (10, 4, 3, 2, 1):
                        # get random cameras
                        cameras = init_random_cameras(
                            cam_type, batch_size, random_z=True
                        ).to(device)
                        # try all alignment modes
                        for mode in ("extrinsics", "centers"):
                            # try different noise levels
                            for add_noise in (0.0, 0.01, 1e-4):
                                self._corresponding_cameras_alignment_test_case(
                                    cameras,
                                    R_align_gt,
                                    T_align_gt,
                                    s_align_gt,
                                    estimate_scale,
                                    mode,
                                    add_noise,
                                )

    def _corresponding_cameras_alignment_test_case(
        self,
        cameras,
        R_align_gt,
        T_align_gt,
        s_align_gt,
        estimate_scale,
        mode,
        add_noise,
    ):
        batch_size = cameras.R.shape[0]

        # get target camera centers
        R_new = torch.bmm(R_align_gt[None].expand_as(cameras.R), cameras.R)
        T_new = (
            torch.bmm(T_align_gt[None, None].repeat(batch_size, 1, 1), cameras.R)[:, 0]
            + cameras.T
        ) * s_align_gt

        if add_noise != 0.0:
            R_new = torch.bmm(R_new, so3_exp_map(torch.randn_like(T_new) * add_noise))
            T_new += torch.randn_like(T_new) * add_noise

        # create new cameras from R_new and T_new
        cameras_tgt = cameras.clone()
        cameras_tgt.R = R_new
        cameras_tgt.T = T_new

        # align cameras and cameras_tgt
        cameras_aligned = corresponding_cameras_alignment(
            cameras, cameras_tgt, estimate_scale=estimate_scale, mode=mode
        )

        if batch_size <= 2 and mode == "centers":
            # underdetermined case - check only the center alignment error
            # since the rotation and translation are ambiguous here
            self.assertClose(
                cameras_aligned.get_camera_center(),
                cameras_tgt.get_camera_center(),
                atol=max(add_noise * 7.0, 1e-4),
            )

        else:

            def _rmse(a):
                return (torch.norm(a, dim=1, p=2) ** 2).mean().sqrt()

            if add_noise != 0.0:
                # in a noisy case check mean rotation/translation error for
                # extrinsic alignment and root mean center error for center alignment
                if mode == "centers":
                    self.assertNormsClose(
                        cameras_aligned.get_camera_center(),
                        cameras_tgt.get_camera_center(),
                        _rmse,
                        atol=max(add_noise * 10.0, 1e-4),
                    )
                elif mode == "extrinsics":
                    angle_err = so3_relative_angle(
                        cameras_aligned.R, cameras_tgt.R, cos_angle=True
                    ).mean()
                    self.assertClose(
                        angle_err, torch.ones_like(angle_err), atol=add_noise * 0.03
                    )
                    self.assertNormsClose(
                        cameras_aligned.T, cameras_tgt.T, _rmse, atol=add_noise * 7.0
                    )
                else:
                    raise ValueError(mode)

            else:
                # compare the rotations and translations of cameras
                self.assertClose(cameras_aligned.R, cameras_tgt.R, atol=3e-4)
                self.assertClose(cameras_aligned.T, cameras_tgt.T, atol=3e-4)
                # compare the centers
                self.assertClose(
                    cameras_aligned.get_camera_center(),
                    cameras_tgt.get_camera_center(),
                    atol=3e-4,
                )

    @staticmethod
    def corresponding_cameras_alignment(
        batch_size: int, estimate_scale: bool, mode: str, cam_type=SfMPerspectiveCameras
    ):
        device = torch.device("cuda:0")
        cameras_src, cameras_tgt = [
            init_random_cameras(cam_type, batch_size, random_z=True).to(device)
            for _ in range(2)
        ]

        torch.cuda.synchronize()

        def compute_corresponding_cameras_alignment():
            corresponding_cameras_alignment(
                cameras_src, cameras_tgt, estimate_scale=estimate_scale, mode=mode
            )
            torch.cuda.synchronize()

        return compute_corresponding_cameras_alignment