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

import platform

import numpy as np
import pytest

from pytorch_layer_test_class import PytorchLayerTest


class TestFloorDivide(PytorchLayerTest):

    def _prepare_input(self):
        return (self.input_tensor, self.other_tensor)

    def create_model(self):
        import torch

        class aten_floor_divide(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, input_tensor, other_tensor):
                return torch.floor_divide(input_tensor, other_tensor)

        return aten_floor_divide(), "aten::floor_divide"

    def create_model_int(self):
        import torch

        class aten_floor_divide(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, input_tensor, other_tensor):
                return torch.floor_divide(input_tensor.to(torch.int32), other_tensor.to(torch.int64))

        return aten_floor_divide(), "aten::floor_divide"

    def create_model_inplace(self):
        import torch

        class aten_floor_divide_(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, input_tensor, other_tensor):
                return input_tensor.floor_divide_(other_tensor), input_tensor

        return aten_floor_divide_(), "aten::floor_divide_"

    @pytest.mark.parametrize('input_tensor', [
        [5], [5, 5, 1], [1, 1, 5, 5],
    ])
    @pytest.mark.parametrize('other_tensor', [
        np.array([[0.5]]).astype(np.float32), [5], [5, 1], [1, 5]
    ])
    @pytest.mark.nightly
    @pytest.mark.precommit
    @pytest.mark.precommit_torch_export
    @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                       reason='Ticket - 122715')
    def test_floor_divide(self, input_tensor, other_tensor, ie_device, precision, ir_version):
        if isinstance(input_tensor, list):
            self.input_tensor = self.random.randn(*input_tensor)
        else:
            self.input_tensor = input_tensor
        if isinstance(other_tensor, list):
            self.other_tensor = self.random.randn(*other_tensor)
        else:
            self.other_tensor = other_tensor
        self._test(*self.create_model(), ie_device, precision, ir_version, trace_model=True, use_convert_model=True)

    @pytest.mark.parametrize('input_tensor', [
        [5, 5, 5], [1, 1, 5, 5],
    ])
    @pytest.mark.parametrize('other_tensor', [
        np.array([0.5]).astype(np.float32), [5], [5, 1], [1, 5]
    ])
    @pytest.mark.nightly
    @pytest.mark.precommit
    @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                       reason='Ticket - 122715')
    def test_floor_divide_(self, input_tensor, other_tensor, ie_device, precision, ir_version):
        if isinstance(input_tensor, list):
            self.input_tensor = self.random.randn(*input_tensor)
        else:
            self.input_tensor = input_tensor
        if isinstance(other_tensor, list):
            self.other_tensor = self.random.randn(*other_tensor)
        else:
            self.other_tensor = other_tensor
        self._test(*self.create_model_inplace(), ie_device, precision, ir_version, trace_model=True, use_convert_model=True)

    @pytest.mark.parametrize('input_data', [
        {"tensor": [5], "low": 0, "high": 10},
        {"tensor": [5, 5, 1], "low": 1, "high": 10},
        {"tensor": [1, 1, 5, 5], "low": 1, "high": 10}
    ])
    @pytest.mark.parametrize('other_data', [
        {"tensor": np.array([[2]]).astype(np.float32)},
        {"tensor": [5], "low": 1, "high": 10},
        {"tensor": [5, 1], "low": 1, "high": 10},
        {"tensor":  [5, 1], "low": 1, "high": 10}
    ])
    @pytest.mark.nightly
    @pytest.mark.precommit
    @pytest.mark.precommit_torch_export
    def test_floor_divide_int(self, input_data, other_data, ie_device, precision, ir_version):
        input_tensor = input_data["tensor"]
        if isinstance(input_tensor, list):
            self.input_tensor = self.random.randint(input_data["low"],
                                                    input_data["high"],
                                                    size=input_tensor, dtype=np.float32)
        else:
            self.input_tensor = input_tensor

        other_tensor = other_data["tensor"]
        if isinstance(other_tensor, list):
            self.other_tensor = self.random.randint(other_data["low"],
                                                    other_data["high"],
                                                    size=other_tensor, dtype=np.float32)
        else:
            self.other_tensor = other_tensor
        self._test(*self.create_model_int(), ie_device, precision, ir_version)
