# Copyright (C) 2018-2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import numpy as np
import pytest
import torch
from pytorch_layer_test_class import PytorchLayerTest

class TestPolar(PytorchLayerTest):
    def _prepare_input(self, input_shape=(1, 1000), dtype=np.float32):
        return (
            self.random.uniform(0, 10, input_shape, dtype=dtype),
            self.random.uniform(-np.pi, np.pi, input_shape, dtype=dtype)
        )

    def create_model(self):
        class PolarModel(torch.nn.Module):
            def forward(self, abs, angle):
                complex_tensor = torch.polar(abs, angle)
                return torch.view_as_real(complex_tensor)
        return PolarModel(), "aten::polar"

    @pytest.mark.parametrize("input_case", [
        (1, 1000),
        (2, 500),
        (5, 200),
        (10, 100),
    ])
    @pytest.mark.parametrize("dtype", [
        np.float32,
        np.float64
    ])
    @pytest.mark.nightly
    @pytest.mark.precommit
    def test_polar(self, input_case, dtype, ie_device, precision, ir_version):
        self._test(*self.create_model(), ie_device, precision, ir_version,
                   kwargs_to_prepare_input={"input_shape": input_case, "dtype": dtype},
                   trace_model=True, use_convert_model=True, dynamic_shapes=False)
