# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import math

import pytest
import torch

import kornia

from testing.base import BaseTester


class TestCam2Pixel(BaseTester):
    def _create_intrinsics(self, batch_size, fx, fy, cx, cy, device, dtype):
        temp = torch.eye(4, device=device, dtype=dtype)
        temp[0, 0], temp[0, 2] = fx, cx
        temp[1, 1], temp[1, 2] = fy, cy
        intrinsics = temp.expand(batch_size, -1, -1)
        return intrinsics

    def _create_intrinsics_inv(self, batch_size, fx, fy, cx, cy, device, dtype):
        temp = torch.eye(4, device=device, dtype=dtype)
        temp[0, 0], temp[0, 2] = 1 / fx, -cx / fx
        temp[1, 1], temp[1, 2] = 1 / fy, -cy / fy
        intrinsics_inv = temp.expand(batch_size, -1, -1)
        return intrinsics_inv

    def _get_samples(self, shape, low, high, device, dtype):
        """Return a tensor having the given shape and whose values are in the range [low, high)"""
        return ((high - low) * torch.rand(shape, device=device, dtype=dtype)) + low

    @pytest.mark.parametrize("batch_size", (1,))
    def test_smoke(self, batch_size, device, dtype):
        H, W = 250, 500
        fx, fy = W, H
        cx, cy = W / 2, H / 2
        eps = 1e-12
        seed = 77
        low, high = -500, 500

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)

        # Setting the projection matrix to the intrinsic matrix for
        # simplicity (i.e. assuming that the RT matrix is an identity matrix)
        proj_mat = intrinsics

        torch.manual_seed(seed)
        cam_coords_src = self._get_samples((batch_size, H, W, 3), low, high, device, dtype)

        pixel_coords_dst = kornia.geometry.camera.cam2pixel(
            cam_coords_src=cam_coords_src, dst_proj_src=proj_mat, eps=eps
        )
        assert pixel_coords_dst.shape == (batch_size, H, W, 2)

    @pytest.mark.parametrize("batch_size", (1, 2, 5))
    def test_consistency(self, batch_size, device, dtype):
        H, W = 250, 500
        fx, fy = W, H
        cx, cy = W / 2, H / 2
        eps = 1e-12
        seed = 77
        low, high = -500, 500

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        intrinsics_inv = self._create_intrinsics_inv(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)

        # Setting the projection matrix to the intrinsic matrix for
        # simplicity (i.e. assuming that the RT matrix is an identity matrix)
        proj_mat = intrinsics

        torch.manual_seed(seed)
        cam_coords_input = self._get_samples((batch_size, H, W, 3), low, high, device, dtype)

        pixel_coords_output = kornia.geometry.camera.cam2pixel(
            cam_coords_src=cam_coords_input, dst_proj_src=proj_mat, eps=eps
        )

        last_ch = torch.ones((batch_size, H, W, 1), device=device, dtype=dtype)
        pixel_coords_concat = torch.cat([pixel_coords_output, last_ch], axis=-1)

        depth = cam_coords_input[..., 2:3].permute(0, 3, 1, 2).contiguous()
        cam_coords_output = kornia.geometry.camera.pixel2cam(
            depth=depth, intrinsics_inv=intrinsics_inv, pixel_coords=pixel_coords_concat
        )

        self.assert_close(cam_coords_output, cam_coords_input, atol=1e-4, rtol=1e-4)

    @pytest.mark.parametrize("batch_size", (1,))
    def test_gradcheck(self, batch_size, device):
        dtype = torch.float64
        H, W = 10, 20
        fx, fy = W, H
        cx, cy = W / 2, H / 2
        eps = 1e-12
        seed = 77
        low, high = -500, 500
        atol, rtol = 1e-5, 1e-3

        # Different tolerances for the below case.
        if (device.type == "cuda") and (dtype == torch.float64):
            atol, rtol = 1e-4, 1e-2

        # If contiguous() is not called, gradcheck fails
        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype).contiguous()

        # Setting the projection matrix to the intrinsic matrix for
        # simplicity (i.e. assuming that the RT matrix is an identity matrix)
        proj_mat = intrinsics

        torch.manual_seed(seed)
        cam_coords_src = self._get_samples((batch_size, H, W, 3), low, high, device, dtype)

        self.gradcheck(kornia.geometry.camera.cam2pixel, (cam_coords_src, proj_mat, eps), atol=atol, rtol=rtol)


class TestPixel2Cam(BaseTester):
    def _create_intrinsics(self, batch_size, fx, fy, cx, cy, device, dtype):
        temp = torch.eye(4, device=device, dtype=dtype)
        temp[0, 0], temp[0, 2] = fx, cx
        temp[1, 1], temp[1, 2] = fy, cy
        intrinsics = temp.expand(batch_size, -1, -1)
        return intrinsics

    def _create_intrinsics_inv(self, batch_size, fx, fy, cx, cy, device, dtype):
        temp = torch.eye(4, device=device, dtype=dtype)
        temp[0, 0], temp[0, 2] = 1 / fx, -cx / fx
        temp[1, 1], temp[1, 2] = 1 / fy, -cy / fy
        intrinsics_inv = temp.expand(batch_size, -1, -1)
        return intrinsics_inv

    def _get_samples(self, shape, low, high, device, dtype):
        """Return a tensor having the given shape and whose values are in the range [low, high)"""
        return ((high - low) * torch.rand(shape, device=device, dtype=dtype)) + low

    @pytest.mark.parametrize("batch_size", (1, 2, 5))
    def test_smoke(self, batch_size, device, dtype):
        H, W = 250, 500
        fx, fy = W, H
        cx, cy = W / 2, H / 2
        seed = 77
        low_1, high_1 = -500, 500
        low_2, high_2 = -(max(W, H) * 3), (max(W, H) * 3)

        torch.manual_seed(seed)
        depth = self._get_samples((batch_size, 1, H, W), low_1, high_1, device, dtype)
        pixel_coords = self._get_samples((batch_size, H, W, 2), low_2, high_2, device, dtype)

        last_ch = torch.ones((batch_size, H, W, 1), device=device, dtype=dtype)
        pixel_coords_input = torch.cat([pixel_coords, last_ch], axis=-1)

        intrinsics_inv = self._create_intrinsics_inv(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)

        output = kornia.geometry.camera.pixel2cam(
            depth=depth, intrinsics_inv=intrinsics_inv, pixel_coords=pixel_coords_input
        )

        assert output.shape == (batch_size, H, W, 3)

    @pytest.mark.parametrize("batch_size", (1, 2, 5))
    def test_consistency(self, batch_size, device, dtype):
        H, W = 250, 500
        fx, fy = W, H
        cx, cy = W / 2, H / 2
        eps = 1e-12
        seed = 77
        low_1, high_1 = -500, 500
        low_2, high_2 = -(max(W, H) * 3), (max(W, H) * 3)

        torch.manual_seed(seed)
        depth = self._get_samples((batch_size, 1, H, W), low_1, high_1, device, dtype)
        pixel_coords = self._get_samples((batch_size, H, W, 2), low_2, high_2, device, dtype)

        last_ch = torch.ones((batch_size, H, W, 1), device=device, dtype=dtype)
        pixel_coords_input = torch.cat([pixel_coords, last_ch], axis=-1)

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        intrinsics_inv = self._create_intrinsics_inv(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)

        cam_coords = kornia.geometry.camera.pixel2cam(
            depth=depth, intrinsics_inv=intrinsics_inv, pixel_coords=pixel_coords_input
        )

        # Setting the projection matrix to the intrinsic matrix for
        # simplicity (i.e. assuming that the RT matrix is an identity matrix)
        proj_mat = intrinsics
        pixel_coords_output = kornia.geometry.camera.cam2pixel(
            cam_coords_src=cam_coords, dst_proj_src=proj_mat, eps=eps
        )
        pixel_coords_concat = torch.cat([pixel_coords_output, last_ch], axis=-1)

        self.assert_close(pixel_coords_concat, pixel_coords_input, atol=1e-4, rtol=1e-4)

    @pytest.mark.parametrize("batch_size", (1,))
    @pytest.mark.slow
    def test_gradcheck(self, batch_size, device):
        dtype = torch.float64
        H, W = 10, 20
        fx, fy = W, H
        cx, cy = W / 2, H / 2
        seed = 77
        low_1, high_1 = -500, 500
        low_2, high_2 = -(max(W, H) * 3), (max(W, H) * 3)

        torch.manual_seed(seed)
        depth = self._get_samples((batch_size, 1, H, W), low_1, high_1, device, dtype)
        pixel_coords = self._get_samples((batch_size, H, W, 2), low_2, high_2, device, dtype)

        last_ch = torch.ones((batch_size, H, W, 1), device=device, dtype=dtype)
        pixel_coords_input = torch.cat([pixel_coords, last_ch], axis=-1)

        # If contiguous() is not called, gradcheck fails
        intrinsics_inv = self._create_intrinsics_inv(
            batch_size, fx, fy, cx, cy, device=device, dtype=dtype
        ).contiguous()

        self.gradcheck(kornia.geometry.camera.pixel2cam, (depth, intrinsics_inv, pixel_coords_input), fast_mode=False)


class TestPinholeCamera(BaseTester):
    def _create_intrinsics(self, batch_size, fx, fy, cx, cy, device, dtype):
        intrinsics = torch.eye(4, device=device, dtype=dtype)
        intrinsics[..., 0, 0] = fx
        intrinsics[..., 1, 1] = fy
        intrinsics[..., 0, 2] = cx
        intrinsics[..., 1, 2] = cy
        return intrinsics.expand(batch_size, -1, -1)

    def _create_extrinsics(self, batch_size, tx, ty, tz, device, dtype):
        extrinsics = torch.eye(4, device=device, dtype=dtype)
        extrinsics[..., 0, -1] = tx
        extrinsics[..., 1, -1] = ty
        extrinsics[..., 2, -1] = tz
        return extrinsics.expand(batch_size, -1, -1)

    def _create_extrinsics_with_rotation(self, batch_size, alpha, beta, gamma, tx, ty, tz, device, dtype):
        Rx = torch.eye(3, device=device, dtype=dtype)
        Rx[1, 1] = math.cos(alpha)
        Rx[1, 2] = math.sin(alpha)
        Rx[2, 1] = -Rx[1, 2]
        Rx[2, 2] = Rx[1, 1]

        Ry = torch.eye(3, device=device, dtype=dtype)
        Ry[0, 0] = math.cos(beta)
        Ry[0, 2] = -math.sin(beta)
        Ry[2, 0] = -Ry[0, 2]
        Ry[2, 2] = Ry[0, 0]

        Rz = torch.eye(3, device=device, dtype=dtype)
        Rz[0, 0] = math.cos(gamma)
        Rz[0, 1] = math.sin(gamma)
        Rz[1, 0] = -Rz[0, 1]
        Rz[1, 1] = Rz[0, 0]

        Ryz = torch.matmul(Ry, Rz)
        R = torch.matmul(Rx, Ryz)

        extrinsics = torch.eye(4, device=device, dtype=dtype)
        extrinsics[..., 0, -1] = tx
        extrinsics[..., 1, -1] = ty
        extrinsics[..., 2, -1] = tz
        extrinsics[:3, :3] = R

        return extrinsics.expand(batch_size, -1, -1)

    def test_smoke(self, device, dtype):
        intrinsics = torch.eye(4, device=device, dtype=dtype)[None]
        extrinsics = torch.eye(4, device=device, dtype=dtype)[None]
        height = torch.ones(1, device=device, dtype=dtype)
        width = torch.ones(1, device=device, dtype=dtype)
        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)
        assert isinstance(pinhole, kornia.geometry.camera.PinholeCamera)

    def test_pinhole_camera_attributes(self, device, dtype):
        batch_size = 1
        height, width = 4, 6
        fx, fy, cx, cy = 1, 2, width / 2, height / 2
        tx, ty, tz = 1, 2, 3

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        extrinsics = self._create_extrinsics(batch_size, tx, ty, tz, device=device, dtype=dtype)
        height = torch.ones(batch_size, device=device, dtype=dtype) * height
        width = torch.ones(batch_size, device=device, dtype=dtype) * width

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)

        assert pinhole.batch_size == batch_size
        assert pinhole.fx.item() == fx
        assert pinhole.fy.item() == fy
        assert pinhole.cx.item() == cx
        assert pinhole.cy.item() == cy
        assert pinhole.tx.item() == tx
        assert pinhole.ty.item() == ty
        assert pinhole.tz.item() == tz
        assert pinhole.height.item() == height
        assert pinhole.width.item() == width
        assert pinhole.rt_matrix.shape == (batch_size, 3, 4)
        assert pinhole.camera_matrix.shape == (batch_size, 3, 3)
        assert pinhole.rotation_matrix.shape == (batch_size, 3, 3)
        assert pinhole.translation_vector.shape == (batch_size, 3, 1)

    def test_pinhole_camera_translation_setters(self, device, dtype):
        batch_size = 1
        height, width = 4, 6
        fx, fy, cx, cy = 1, 2, width / 2, height / 2
        tx, ty, tz = 1, 2, 3

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        extrinsics = self._create_extrinsics(batch_size, tx, ty, tz, device=device, dtype=dtype)
        height = torch.ones(batch_size, device=device, dtype=dtype) * height
        width = torch.ones(batch_size, device=device, dtype=dtype) * width

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)

        assert pinhole.tx.item() == tx
        assert pinhole.ty.item() == ty
        assert pinhole.tz.item() == tz

        # add offset
        pinhole.tx += 3.0
        pinhole.ty += 2.0
        pinhole.tz += 1.0

        assert pinhole.tx.item() == tx + 3.0
        assert pinhole.ty.item() == ty + 2.0
        assert pinhole.tz.item() == tz + 1.0

        # set to zero
        pinhole.tx = 0.0
        pinhole.ty = 0.0
        pinhole.tz = 0.0

        assert pinhole.tx.item() == 0.0
        assert pinhole.ty.item() == 0.0
        assert pinhole.tz.item() == 0.0

    def test_pinhole_camera_attributes_batch2(self, device, dtype):
        batch_size = 2
        height, width = 4, 6
        fx, fy, cx, cy = 1, 2, width / 2, height / 2
        tx, ty, tz = 1, 2, 3

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        extrinsics = self._create_extrinsics(batch_size, tx, ty, tz, device=device, dtype=dtype)
        height = torch.ones(batch_size, device=device, dtype=dtype) * height
        width = torch.ones(batch_size, device=device, dtype=dtype) * width

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)

        assert pinhole.batch_size == batch_size
        assert pinhole.fx.shape[0] == batch_size
        assert pinhole.fy.shape[0] == batch_size
        assert pinhole.cx.shape[0] == batch_size
        assert pinhole.cy.shape[0] == batch_size
        assert pinhole.tx.shape[0] == batch_size
        assert pinhole.ty.shape[0] == batch_size
        assert pinhole.tz.shape[0] == batch_size
        assert pinhole.height.shape[0] == batch_size
        assert pinhole.width.shape[0] == batch_size
        assert pinhole.rt_matrix.shape == (batch_size, 3, 4)
        assert pinhole.camera_matrix.shape == (batch_size, 3, 3)
        assert pinhole.rotation_matrix.shape == (batch_size, 3, 3)
        assert pinhole.translation_vector.shape == (batch_size, 3, 1)

    def test_pinhole_camera_scale(self, device, dtype):
        batch_size = 2
        height, width = 4, 6
        fx, fy, cx, cy = 1, 2, width / 2, height / 2
        tx, ty, tz = 1, 2, 3
        scale_val = 2.0

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        extrinsics = self._create_extrinsics(batch_size, tx, ty, tz, device=device, dtype=dtype)
        height = torch.ones(batch_size, device=device, dtype=dtype) * height
        width = torch.ones(batch_size, device=device, dtype=dtype) * width
        scale_factor = torch.ones(batch_size, device=device, dtype=dtype) * scale_val

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)
        pinhole_scale = pinhole.scale(scale_factor)

        self.assert_close(
            pinhole_scale.intrinsics[..., 0, 0], pinhole.intrinsics[..., 0, 0] * scale_val, atol=1e-4, rtol=1e-4
        )  # fx
        self.assert_close(
            pinhole_scale.intrinsics[..., 1, 1], pinhole.intrinsics[..., 1, 1] * scale_val, atol=1e-4, rtol=1e-4
        )  # fy
        self.assert_close(
            pinhole_scale.intrinsics[..., 0, 2], pinhole.intrinsics[..., 0, 2] * scale_val, atol=1e-4, rtol=1e-4
        )  # cx
        self.assert_close(
            pinhole_scale.intrinsics[..., 1, 2], pinhole.intrinsics[..., 1, 2] * scale_val, atol=1e-4, rtol=1e-4
        )  # cy
        self.assert_close(pinhole_scale.height, pinhole.height * scale_val, atol=1e-4, rtol=1e-4)
        self.assert_close(pinhole_scale.width, pinhole.width * scale_val, atol=1e-4, rtol=1e-4)

    def test_pinhole_camera_scale_inplace(self, device, dtype):
        batch_size = 2
        height, width = 4, 6
        fx, fy, cx, cy = 1, 2, width / 2, height / 2
        tx, ty, tz = 1, 2, 3
        scale_val = 2.0

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        extrinsics = self._create_extrinsics(batch_size, tx, ty, tz, device=device, dtype=dtype)
        height = torch.ones(batch_size, device=device, dtype=dtype) * height
        width = torch.ones(batch_size, device=device, dtype=dtype) * width
        scale_factor = torch.ones(batch_size, device=device, dtype=dtype) * scale_val

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)
        pinhole_scale = pinhole.clone()
        pinhole_scale.scale_(scale_factor)

        self.assert_close(
            pinhole_scale.intrinsics[..., 0, 0], pinhole.intrinsics[..., 0, 0] * scale_val, atol=1e-4, rtol=1e-4
        )  # fx
        self.assert_close(
            pinhole_scale.intrinsics[..., 1, 1], pinhole.intrinsics[..., 1, 1] * scale_val, atol=1e-4, rtol=1e-4
        )  # fy
        self.assert_close(
            pinhole_scale.intrinsics[..., 0, 2], pinhole.intrinsics[..., 0, 2] * scale_val, atol=1e-4, rtol=1e-4
        )  # cx
        self.assert_close(
            pinhole_scale.intrinsics[..., 1, 2], pinhole.intrinsics[..., 1, 2] * scale_val, atol=1e-4, rtol=1e-4
        )  # cy
        self.assert_close(pinhole_scale.height, pinhole.height * scale_val, atol=1e-4, rtol=1e-4)
        self.assert_close(pinhole_scale.width, pinhole.width * scale_val, atol=1e-4, rtol=1e-4)

    def test_pinhole_camera_project_and_unproject(self, device, dtype):
        batch_size = 5
        n = 2  # Point per batch
        height, width = 4, 6
        fx, fy, cx, cy = 1, 2, width / 2, height / 2
        alpha, beta, gamma = 0.0, 0.0, 0.4
        tx, ty, tz = 0, 0, 3

        intrinsics = self._create_intrinsics(batch_size, fx, fy, cx, cy, device=device, dtype=dtype)
        extrinsics = self._create_extrinsics_with_rotation(
            batch_size, alpha, beta, gamma, tx, ty, tz, device=device, dtype=dtype
        )

        height = torch.ones(batch_size, device=device, dtype=dtype) * height
        width = torch.ones(batch_size, device=device, dtype=dtype) * width

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)

        point_3d = torch.rand((batch_size, n, 3), device=device, dtype=dtype)

        depth = point_3d[..., -1:] + tz

        point_2d = pinhole.project(point_3d)
        point_3d_hat = pinhole.unproject(point_2d, depth)
        self.assert_close(point_3d, point_3d_hat, atol=1e-4, rtol=1e-4)

    def test_pinhole_camera_device(self, device, dtype):
        batch_size = 5
        intrinsics = torch.rand((batch_size, 4, 4), device=device, dtype=dtype)
        extrinsics = torch.rand((batch_size, 4, 4), device=device, dtype=dtype)
        height = torch.randint(low=5, high=9, size=(batch_size,), device=device)
        width = torch.randint(low=5, high=9, size=(batch_size,), device=device)

        pinhole = kornia.geometry.camera.PinholeCamera(intrinsics, extrinsics, height, width)
        assert pinhole.device() == intrinsics.device
