# 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.
"""
Tests for Example NPU Backend

This test file demonstrates how to test a custom NPU backend
implementation using TVM's testing infrastructure.
"""

import numpy as np
import pytest

import tvm
import tvm.testing
from tvm import relax
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
from tvm.script import relax as R


@tvm.script.ir_module
class MatmulReLU:
    """Example module with matrix multiplication and ReLU"""

    @R.function
    def main(
        x: R.Tensor((2, 4), "float32"),
        w: R.Tensor((4, 8), "float32"),
    ) -> R.Tensor((2, 8), "float32"):
        with R.dataflow():
            y = relax.op.matmul(x, w)
            z = relax.op.nn.relu(y)
            R.output(z)
        return z


@tvm.script.ir_module
class Conv2dReLU:
    """Example module with 2D convolution and ReLU"""

    @R.function
    def main(
        x: R.Tensor((1, 3, 32, 32), "float32"),
        w: R.Tensor((16, 3, 3, 3), "float32"),
    ) -> R.Tensor((1, 16, 30, 30), "float32"):
        with R.dataflow():
            y = relax.op.nn.conv2d(x, w)
            z = relax.op.nn.relu(y)
            R.output(z)
        return z


@tvm.script.ir_module
class MultipleOps:
    """Example module with multiple operations that can be fused"""

    @R.function
    def main(
        x: R.Tensor((1, 16, 32, 32), "float32"),
    ) -> R.Tensor((1, 16, 16, 16), "float32"):
        with R.dataflow():
            # First ReLU
            y = relax.op.nn.relu(x)
            # Max pooling
            z = relax.op.nn.max_pool2d(y, pool_size=(2, 2), strides=(2, 2))
            # Second ReLU
            out = relax.op.nn.relu(z)
            R.output(out)
        return out


@tvm.script.ir_module
class Softmax:
    """Example module with softmax"""

    @R.function
    def main(x: R.Tensor((2, 8), "float32")) -> R.Tensor((2, 8), "float32"):
        with R.dataflow():
            z = relax.op.nn.softmax(x)
            R.output(z)
        return z


# Check if the example NPU runtime is available
has_example_npu_codegen = tvm.get_global_func("relax.ext.example_npu", True)
has_example_npu_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
has_example_npu = has_example_npu_codegen and has_example_npu_runtime

example_npu_enabled = pytest.mark.skipif(
    not has_example_npu,
    reason="Example NPU backend not enabled. Compile with the example NPU runtime.",
)


def test_example_npu_patterns_registered():
    """Test that all expected patterns are registered"""
    import tvm.relax.backend.contrib.example_npu  # noqa: F401

    patterns = get_patterns_with_prefix("example_npu")
    pattern_names = {p.name for p in patterns}

    # Core patterns that should always be available
    core_patterns = {
        "example_npu.dense",
        "example_npu.matmul",
        "example_npu.conv1d",
        "example_npu.conv2d",
        "example_npu.max_pool2d",
    }

    assert core_patterns.issubset(pattern_names), (
        f"Missing core patterns: {core_patterns - pattern_names}"
    )

    # Check that at least some activation patterns are available
    activation_patterns = {name for name in pattern_names if "relu" in name or "sigmoid" in name}
    assert len(activation_patterns) > 0, "No activation patterns found"


@example_npu_enabled
def test_example_npu_matmul_relu_partitioning():
    """Test graph partitioning for MatMul + ReLU pattern"""
    import tvm.relax.backend.contrib.example_npu  # noqa: F401

    mod = MatmulReLU
    patterns = get_patterns_with_prefix("example_npu")

    # Partition the graph
    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)

    # Verify partitioning happened
    assert partitioned_mod is not None

    # Check that composite functions were created
    for gvar, func in partitioned_mod.functions.items():
        if gvar.name_hint != "main":
            # This should be a composite function
            assert "Composite" in str(func)


@example_npu_enabled
def test_example_npu_conv2d_relu_partitioning():
    """Test graph partitioning for Conv2D + ReLU pattern"""
    import tvm.relax.backend.contrib.example_npu  # noqa: F401

    mod = Conv2dReLU
    patterns = get_patterns_with_prefix("example_npu")

    # Partition the graph
    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)

    assert partitioned_mod is not None


@example_npu_enabled
def test_example_npu_multiple_ops():
    """Test partitioning with multiple fusable operations"""
    import tvm.relax.backend.contrib.example_npu  # noqa: F401

    mod = MultipleOps
    patterns = get_patterns_with_prefix("example_npu")

    # Partition the graph
    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)

    assert partitioned_mod is not None


@example_npu_enabled
def test_example_npu_softmax_partitioning():
    """Test graph partitioning for softmax pattern"""
    import tvm.relax.backend.contrib.example_npu  # noqa: F401

    mod = Softmax
    patterns = get_patterns_with_prefix("example_npu")

    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)

    assert partitioned_mod is not None

    for gvar, func in partitioned_mod.functions.items():
        if gvar.name_hint != "main":
            assert "Composite" in str(func)


@example_npu_enabled
def test_example_npu_codegen():
    """Test code generation for the example NPU backend"""
    import tvm.relax.backend.contrib.example_npu  # noqa: F401

    mod = MatmulReLU
    patterns = get_patterns_with_prefix("example_npu")

    # Partition and generate code
    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
    partitioned_mod = RunCodegen()(partitioned_mod)

    assert partitioned_mod is not None

    # The module should now contain external function calls
    main_func = partitioned_mod["main"]
    assert main_func is not None


@example_npu_enabled
def test_example_npu_runtime_execution():
    """Test end-to-end execution with the example NPU runtime"""
    import tvm.relax.backend.contrib.example_npu

    # Create simple test inputs
    np.random.seed(42)
    x_np = np.random.randn(2, 4).astype("float32")
    w_np = np.random.randn(4, 8).astype("float32")

    # Expected output (computed with NumPy)
    expected = np.maximum(0, np.matmul(x_np, w_np))

    # Build and run with example NPU backend
    mod = MatmulReLU
    patterns = get_patterns_with_prefix("example_npu")

    # Apply transformations
    mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
    mod = MergeCompositeFunctions()(mod)
    mod = RunCodegen()(mod)

    # Build the module
    target = tvm.target.Target("llvm")
    with tvm.transform.PassContext(opt_level=3):
        built = relax.build(mod, target)

    # Create VM and run
    vm = relax.VirtualMachine(built, tvm.cpu())

    x_tvm = tvm.runtime.tensor(x_np, tvm.cpu())
    w_tvm = tvm.runtime.tensor(w_np, tvm.cpu())

    result = vm["main"](x_tvm, w_tvm)

    # Verify the result shape is correct (the runtime is a stub and does not compute numerically)
    assert result.numpy().shape == expected.shape


if __name__ == "__main__":
    # Run tests locally for debugging
    test_example_npu_patterns_registered()

    if has_example_npu:
        print("Example NPU backend is available, running tests...")
        test_example_npu_matmul_relu_partitioning()
        test_example_npu_conv2d_relu_partitioning()
        test_example_npu_softmax_partitioning()
        test_example_npu_multiple_ops()
        test_example_npu_codegen()
        test_example_npu_runtime_execution()
        print("All tests passed!")
    else:
        print("Example NPU backend not available. Compile with example NPU runtime to run tests.")