# 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.
# pylint: disable=invalid-name
"""Testing utilities for relax VM"""

from typing import Any

import numpy as np  # type: ignore

import tvm
from tvm import relax
from tvm.runtime import Object


@tvm.register_global_func("test.vm.move")
def move(src):
    return src


@tvm.register_global_func("test.vm.add")
def add(a, b):
    ret = a.numpy() + b.numpy()
    return tvm.runtime.tensor(ret)


@tvm.register_global_func("test.vm.mul")
def mul(a, b):
    ret = a.numpy() * b.numpy()
    return tvm.runtime.tensor(ret)


@tvm.register_global_func("test.vm.equal_zero")
def equal_zero(a):
    ret = np.all(a.numpy() == 0)
    return tvm.runtime.tensor(ret)


@tvm.register_global_func("test.vm.subtract_one")
def subtract_one(a):
    ret = np.subtract(a.numpy(), 1)
    return tvm.runtime.tensor(ret)


@tvm.register_global_func("test.vm.identity")
def identity_packed(a, b):
    b[:] = tvm.runtime.tensor(a.numpy())


@tvm.register_global_func("test.vm.tile")
def tile_packed(a, b):
    b[:] = tvm.runtime.tensor(np.tile(a.numpy(), (1, 2)))


@tvm.register_global_func("test.vm.add_scalar")
def add_scalar(a, b):
    return a + b


@tvm.register_global_func("test.vm.get_device_id")
def get_device_id(device):
    return device.index


def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any]) -> Object:
    # uses save_function to create a closure with the given inputs
    # and ensure the result is the same
    # (assumes the functions return tensors and that they're idempotent)
    saved_name = f"{func_name}_saved"
    vm.save_function(func_name, saved_name, *inputs)
    res1 = vm[func_name](*inputs)
    res2 = vm[saved_name]()
    tvm.testing.assert_allclose(res1.numpy(), res2.numpy(), rtol=1e-7, atol=1e-7)
    return res1


@tvm.register_global_func("test.vm.check_if_defined")
def check_if_defined(obj: tvm.Object) -> tvm.tirx.IntImm:
    return tvm.runtime.convert(obj is not None)