import pytest
import torch

import kornia
from kornia.testing import assert_close


class TestRenderGaussian2d:
    @pytest.fixture
    def gaussian(self, device, dtype):
        return torch.tensor(
            [
                [0.002969, 0.013306, 0.021938, 0.013306, 0.002969],
                [0.013306, 0.059634, 0.098320, 0.059634, 0.013306],
                [0.021938, 0.098320, 0.162103, 0.098320, 0.021938],
                [0.013306, 0.059634, 0.098320, 0.059634, 0.013306],
                [0.002969, 0.013306, 0.021938, 0.013306, 0.002969],
            ],
            dtype=dtype,
            device=device,
        )

    def test_pixel_coordinates(self, gaussian, device, dtype):
        mean = torch.tensor([2.0, 2.0], dtype=dtype, device=device)
        std = torch.tensor([1.0, 1.0], dtype=dtype, device=device)
        actual = kornia.geometry.subpix.render_gaussian2d(mean, std, (5, 5), False)
        assert_close(actual, gaussian, rtol=0, atol=1e-4)

    def test_normalized_coordinates(self, gaussian, device, dtype):
        mean = torch.tensor([0.0, 0.0], dtype=dtype, device=device)
        std = torch.tensor([0.25, 0.25], dtype=dtype, device=device)
        actual = kornia.geometry.subpix.render_gaussian2d(mean, std, (5, 5), True)
        assert_close(actual, gaussian, rtol=0, atol=1e-4)

    def test_jit(self, device, dtype):
        mean = torch.tensor([0.0, 0.0], dtype=dtype, device=device)
        std = torch.tensor([0.25, 0.25], dtype=dtype, device=device)
        args = (mean, std, (5, 5), True)
        op = kornia.geometry.subpix.render_gaussian2d
        op_jit = torch.jit.script(op)
        assert_close(op(*args), op_jit(*args), rtol=0, atol=1e-5)

    @pytest.mark.skip(reason="it works but raises some warnings.")
    def test_jit_trace(self, device, dtype):
        def op(mean, std):
            return kornia.geometry.subpix.render_gaussian2d(mean, std, (5, 5), True)

        mean = torch.tensor([0.0, 0.0], dtype=dtype, device=device)
        std = torch.tensor([0.25, 0.25], dtype=dtype, device=device)
        args = (mean, std)
        op_jit = torch.jit.trace(op, args)
        assert_close(op(*args), op_jit(*args), rtol=0, atol=1e-5)


class TestSpatialSoftmax2d:
    @pytest.fixture(params=[torch.ones(1, 1, 5, 7), torch.randn(2, 3, 16, 16)])
    def input(self, request, device, dtype):
        return request.param.to(device, dtype)

    def test_forward(self, input):
        actual = kornia.geometry.subpix.spatial_softmax2d(input)
        assert actual.lt(0).sum().item() == 0, 'expected no negative values'
        sums = actual.sum(-1).sum(-1)
        assert_close(sums, torch.ones_like(sums))

    def test_jit(self, input):
        op = kornia.geometry.subpix.spatial_softmax2d
        op_jit = torch.jit.script(op)
        assert_close(op(input), op_jit(input), rtol=0, atol=1e-5)

    @pytest.mark.skip(reason="it works but raises some warnings.")
    def test_jit_trace(self, input):
        op = kornia.geometry.subpix.spatial_softmax2d
        op_jit = torch.jit.trace(op, (input,))
        assert_close(op(input), op_jit(input), rtol=0, atol=1e-5)


class TestSpatialExpectation2d:
    @pytest.fixture(
        params=[
            (
                torch.tensor([[[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0]]]]),
                torch.tensor([[[1.0, -1.0]]]),
                torch.tensor([[[2.0, 0.0]]]),
            )
        ]
    )
    def example(self, request, device, dtype):
        input, expected_norm, expected_px = request.param
        return input.to(device, dtype), expected_norm.to(device, dtype), expected_px.to(device, dtype)

    def test_forward(self, example):
        input, expected_norm, expected_px = example
        actual_norm = kornia.geometry.subpix.spatial_expectation2d(input, True)
        assert_close(actual_norm, expected_norm)
        actual_px = kornia.geometry.subpix.spatial_expectation2d(input, False)
        assert_close(actual_px, expected_px)

    def test_jit(self, example):
        input = example[0]
        op = kornia.geometry.subpix.spatial_expectation2d
        op_jit = torch.jit.script(op)
        assert_close(op(input), op_jit(input), rtol=0, atol=1e-5)

    @pytest.mark.skip(reason="it works but raises some warnings.")
    def test_jit_trace(self, example):
        input = example[0]
        op = kornia.geometry.subpix.spatial_expectation2d
        op_jit = torch.jit.trace(op, (input,))
        assert_close(op(input), op_jit(input), rtol=0, atol=1e-5)