# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# ruff: noqa: F401
"""
Test BasePyModule core functionality.

This test verifies:
1. BasePyModule instantiation and basic methods
2. TIR function compilation and execution
3. Python function integration
4. DLPack conversion between PyTorch and TVM
"""

import numpy as np
import pytest
import torch

import tvm
from tvm import relax, tirx
from tvm.relax import BasePyModule
from tvm.script import relax as R
from tvm.script import tirx as T


class TestBasePyModule:
    """Test BasePyModule core functionality."""

    def test_base_py_module_instantiation(self):
        @T.prim_func(s_tir=True)
        def simple_func(A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32")):
            for i in T.grid(10):
                B[i] = A[i] * 2.0

        ir_mod = tvm.IRModule({"simple_func": simple_func})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        assert isinstance(py_mod, BasePyModule)
        assert hasattr(py_mod, "call_tir")
        assert hasattr(py_mod, "call_dps_packed")
        assert hasattr(py_mod, "compiled_tir_funcs")

    def test_base_py_module_instantiation_gpu(self):
        @T.prim_func(s_tir=True)
        def simple_func(A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32")):
            for i in T.grid(10):
                B[i] = A[i] * 2.0

        ir_mod = tvm.IRModule({"simple_func": simple_func})

        if tvm.cuda().exist:

            def run_and_check():
                device = tvm.cuda(0)
                py_mod = BasePyModule(ir_mod, device)

                assert isinstance(py_mod, BasePyModule)
                assert hasattr(py_mod, "call_tir")
                assert hasattr(py_mod, "call_dps_packed")
                assert hasattr(py_mod, "compiled_tir_funcs")
                # Check if target contains "cuda" instead of exact match
                assert "cuda" in str(py_mod.target)

            tvm.testing.run_with_gpu_lock(run_and_check)
        else:
            pytest.skip("CUDA not available")

    def test_tir_function_compilation(self):
        @T.prim_func(s_tir=True)
        def add_func(
            A: T.Buffer((5,), "float32"), B: T.Buffer((5,), "float32"), C: T.Buffer((5,), "float32")
        ):
            for i in T.grid(5):
                C[i] = A[i] + B[i]

        ir_mod = tvm.IRModule({"add_func": add_func})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        assert "add_func" in py_mod.tir_func_names
        assert "add_func" in py_mod.compiled_tir_funcs

    def test_call_tir_with_pytorch_tensors(self):
        @T.prim_func(s_tir=True)
        def scale_func(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
            for i in T.grid(4):
                B[i] = A[i] * T.float32(2.5)

        ir_mod = tvm.IRModule({"scale_func": scale_func})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
        scale_value = 2.5

        result = py_mod.call_tir(scale_func, [input_tensor], R.Tensor((4,), "float32"))

        assert isinstance(result, torch.Tensor)
        assert result.shape == (4,)
        expected = input_tensor * scale_value
        assert torch.allclose(result, expected, atol=1e-5)

    def test_call_tir_with_pytorch_tensors_gpu(self):
        if tvm.cuda().exist:

            def run_and_check():
                # Create a simple IRModule without TIR functions for GPU testing
                ir_mod = tvm.IRModule({})
                device = tvm.cuda(0)
                py_mod = BasePyModule(ir_mod, device)

                # Test basic GPU functionality without TIR compilation issues
                assert isinstance(py_mod, BasePyModule)
                assert hasattr(py_mod, "call_tir")
                assert hasattr(py_mod, "call_dps_packed")
                assert "cuda" in str(py_mod.target)

                # Test that we can create GPU tensors and they work
                input_tensor = torch.tensor(
                    [1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device="cuda"
                )
                assert input_tensor.device.type == "cuda"
                assert input_tensor.shape == (4,)

            tvm.testing.run_with_gpu_lock(run_and_check)
        else:
            pytest.skip("CUDA not available")

    def test_dlpack_conversion_pytorch_to_tvm(self):
        @T.prim_func(s_tir=True)
        def identity_func(A: T.Buffer((3,), "float32"), B: T.Buffer((3,), "float32")):
            for i in T.grid(3):
                B[i] = A[i]

        ir_mod = tvm.IRModule({"identity_func": identity_func})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)

        result = py_mod.call_tir(identity_func, [input_tensor], R.Tensor((3,), "float32"))

        assert isinstance(result, torch.Tensor)
        assert torch.allclose(result, input_tensor, atol=1e-5)

    def test_dlpack_conversion_tvm_to_pytorch(self):
        @T.prim_func(s_tir=True)
        def constant_func(B: T.Buffer((2,), "float32")):
            for i in T.grid(2):
                B[i] = T.float32(5.0)

        ir_mod = tvm.IRModule({"constant_func": constant_func})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        result = py_mod.call_tir(constant_func, [], R.Tensor((2,), "float32"))

        assert isinstance(result, torch.Tensor)
        assert result.shape == (2,)
        expected = torch.tensor([5.0, 5.0], dtype=torch.float32)
        assert torch.allclose(result, expected, atol=1e-5)

    def test_add_python_function(self):
        ir_mod = tvm.IRModule({})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        def custom_activation(x):
            return torch.tanh(x)

        py_mod.add_python_function("custom_activation", custom_activation)

        assert hasattr(py_mod, "custom_activation")
        assert "custom_activation" in py_mod.pyfuncs

        input_tensor = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float32)
        result = py_mod.custom_activation(input_tensor)

        assert isinstance(result, torch.Tensor)
        expected = torch.tanh(input_tensor)
        assert torch.allclose(result, expected, atol=1e-5)

    def test_call_dps_packed_with_python_function(self):
        ir_mod = tvm.IRModule({})
        device = tvm.cpu(0)
        py_mod = BasePyModule(ir_mod, device)

        def my_softmax(tensor, dim):
            return torch.softmax(tensor, dim=dim)

        py_mod.add_python_function("my_softmax", my_softmax)

        input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)

        result = py_mod.call_dps_packed(
            "my_softmax", [input_tensor, 1], R.Tensor((2, 2), "float32")
        )

        assert isinstance(result, torch.Tensor)
        expected = torch.softmax(input_tensor, dim=1)
        assert torch.allclose(result, expected, atol=1e-5)


if __name__ == "__main__":
    tvm.testing.main()
