# Copyright 2019 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import json
import numpy as np
import unittest

from iree import runtime as rt
from iree.runtime.function import (
    FunctionInvoker,
    IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
    IMPLICIT_BUFFER_ARG_USAGE,
)
from iree.runtime._binding import VmVariantList


class MockVmContext:

  def __init__(self, invoke_callback):
    self._invoke_callback = invoke_callback
    self.invocations = []

  def invoke(self, vm_function, arg_list, ret_list):
    self._invoke_callback(arg_list, ret_list)
    self.invocations.append((vm_function, arg_list, ret_list))
    print(f"INVOKE: {arg_list} -> {ret_list}")

  @property
  def mock_arg_reprs(self):
    return repr([arg_list for _, arg_list, _ in self.invocations])


class MockVmFunction:

  def __init__(self, reflection):
    self.reflection = reflection


class FunctionTest(unittest.TestCase):

  @classmethod
  def setUpClass(cls):
    # Doesn't matter what device. We just need one.
    config = rt.Config("local-task")
    cls.device = config.device

  def testNoReflectionScalars(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)
      ret_list.push_int(4)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker(1, 2)
    self.assertEqual("[<VmVariantList(2): [1, 2]>]", vm_context.mock_arg_reprs)
    self.assertEqual((3, 4), result)

  def testKeywordArgs(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [
                        "i32",
                        ["named", "a", "i32"],
                        ["named", "b", "i32"],
                    ],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker(-1, a=1, b=2)
    self.assertEqual("[<VmVariantList(3): [-1, 1, 2]>]",
                     vm_context.mock_arg_reprs)
    self.assertEqual(3, result)

  def testListArg(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi":
            json.dumps({
                "a": [["slist", "i32", "i32"],],
                "r": ["i32",],
            })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker([2, 3])
    self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]",
                     vm_context.mock_arg_reprs)

  def testListArgNoReflection(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker([2, 3])
    self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]",
                     vm_context.mock_arg_reprs)

  def testListArgArityMismatch(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi":
            json.dumps({
                "a": [["slist", "i32", "i32"],],
                "r": ["i32",],
            })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError,
                                "expected a sequence with 2 values. got:"):
      _ = invoker([2, 3, 4])

  def testTupleArg(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi":
            json.dumps({
                "a": [["stuple", "i32", "i32"],],
                "r": ["i32",],
            })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker((2, 3))
    self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]",
                     vm_context.mock_arg_reprs)

  def testDictArg(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [["sdict", ["a", "i32"], ["b", "i32"]],],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker({"b": 3, "a": 2})
    self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]",
                     vm_context.mock_arg_reprs)

  def testDictArgArityMismatch(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [["sdict", ["a", "i32"], ["b", "i32"]],],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError,
                                "expected a dict with 2 values. got:"):
      _ = invoker({"a": 2, "b": 3, "c": 4})

  def testDictArgKeyError(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [["sdict", ["a", "i32"], ["b", "i32"]],],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError, "could not get item 'b' from: "):
      _ = invoker({"a": 2, "c": 3})

  def testDictArgNoReflection(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker({"b": 3, "a": 2})
    self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]",
                     vm_context.mock_arg_reprs)

  def testInlinedResults(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)
      ret_list.push_int(4)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi": json.dumps({
            "a": [],
            "r": [["slist", "i32", "i32"]],
        })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker()
    self.assertEqual([3, 4], result)

  def testNestedResults(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)
      sub_list = VmVariantList(2)
      sub_dict = VmVariantList(2)
      sub_dict.push_int(100)
      sub_dict.push_int(200)
      sub_list.push_list(sub_dict)
      sub_list.push_int(6)
      ret_list.push_list(sub_list)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [],
                    "r": [
                        "i32",
                        [
                            "slist",
                            ["sdict", ["bar", "i32"], ["foo", "i32"]],
                            "i64",
                        ]
                    ],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker()
    self.assertEqual((3, [{"bar": 100, "foo": 200}, 6]), result)

  def testMissingPositional(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [
                        "i32",
                        ["named", "a", "i32"],
                        ["named", "b", "i32"],
                    ],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError, "mismatched call arity:"):
      result = invoker(a=1, b=1)

  def testMissingPositionalNdarray(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [
                        ["ndarray", "i32", 1, 1],
                        ["named", "a", ["ndarray", "i32", 1, 1]],
                        ["named", "b", ["ndarray", "i32", 1, 1]],
                    ],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError, "mismatched call arity:"):
      result = invoker(a=1, b=1)

  def testMissingKeyword(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [
                        "i32",
                        ["named", "a", "i32"],
                        ["named", "b", "i32"],
                    ],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError, "mismatched call arity:"):
      result = invoker(-1, a=1)

  def testMissingKeywordNdArray(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [
                        ["ndarray", "i32", 1, 1],
                        ["named", "a", ["ndarray", "i32", 1, 1]],
                        ["named", "b", ["ndarray", "i32", 1, 1]],
                    ],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError, "mismatched call arity:"):
      result = invoker(-1, a=1)

  def testExtraKeyword(self):

    def invoke(arg_list, ret_list):
      ret_list.push_int(3)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(
        reflection={
            "iree.abi":
                json.dumps({
                    "a": [
                        "i32",
                        ["named", "a", "i32"],
                        ["named", "b", "i32"],
                    ],
                    "r": ["i32",],
                })
        })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    with self.assertRaisesRegex(ValueError, "specified kwarg 'c' is unknown"):
      result = invoker(-1, a=1, b=2, c=3)

  def testNdarrayArg(self):
    arg_array = np.asarray([1, 0], dtype=np.int32)

    invoked_arg_list = None

    def invoke(arg_list, ret_list):
      nonlocal invoked_arg_list
      invoked_arg_list = arg_list

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi": json.dumps({
            "a": [["ndarray", "i32", 1, 2]],
            "r": [],
        })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker(arg_array)
    self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>",
                     repr(invoked_arg_list))

  def testDeviceArrayArg(self):
    # Note that since the device array is set up to disallow implicit host
    # transfers, this also verifies that no accidental/automatic transfers
    # are done as part of marshalling the array to the function.
    arg_array = rt.asdevicearray(self.device,
                                 np.asarray([1, 0], dtype=np.int32),
                                 implicit_host_transfer=False)

    invoked_arg_list = None

    def invoke(arg_list, ret_list):
      nonlocal invoked_arg_list
      invoked_arg_list = arg_list

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi": json.dumps({
            "a": [["ndarray", "i32", 1, 2]],
            "r": [],
        })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker(arg_array)
    self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>",
                     repr(invoked_arg_list))

  def testBufferViewArg(self):
    arg_buffer_view = self.device.allocator.allocate_buffer_copy(
        memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
        allowed_usage=IMPLICIT_BUFFER_ARG_USAGE,
        buffer=np.asarray([1, 0], dtype=np.int32),
        element_type=rt.HalElementType.SINT_32)

    invoked_arg_list = None

    def invoke(arg_list, ret_list):
      nonlocal invoked_arg_list
      invoked_arg_list = arg_list

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi": json.dumps({
            "a": [["ndarray", "i32", 1, 2]],
            "r": [],
        })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker(arg_buffer_view)
    self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>",
                     repr(invoked_arg_list))

  def testNdarrayArgNoReflection(self):
    arg_array = np.asarray([1, 0], dtype=np.int32)

    invoked_arg_list = None

    def invoke(arg_list, ret_list):
      nonlocal invoked_arg_list
      invoked_arg_list = arg_list

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker(arg_array)
    self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>",
                     repr(invoked_arg_list))

  def testDeviceArrayArgNoReflection(self):
    # Note that since the device array is set up to disallow implicit host
    # transfers, this also verifies that no accidental/automatic transfers
    # are done as part of marshalling the array to the function.
    arg_array = rt.asdevicearray(self.device,
                                 np.asarray([1, 0], dtype=np.int32),
                                 implicit_host_transfer=False)

    invoked_arg_list = None

    def invoke(arg_list, ret_list):
      nonlocal invoked_arg_list
      invoked_arg_list = arg_list

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker(arg_array)
    self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>",
                     repr(invoked_arg_list))

  def testBufferViewArgNoReflection(self):
    arg_buffer_view = self.device.allocator.allocate_buffer_copy(
        memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
        allowed_usage=IMPLICIT_BUFFER_ARG_USAGE,
        buffer=np.asarray([1, 0], dtype=np.int32),
        element_type=rt.HalElementType.SINT_32)

    invoked_arg_list = None

    def invoke(arg_list, ret_list):
      nonlocal invoked_arg_list
      invoked_arg_list = arg_list

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    _ = invoker(arg_buffer_view)
    self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>",
                     repr(invoked_arg_list))

  def testReturnBufferView(self):
    result_array = np.asarray([1, 0], dtype=np.int32)

    def invoke(arg_list, ret_list):
      buffer_view = self.device.allocator.allocate_buffer_copy(
          memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
          allowed_usage=IMPLICIT_BUFFER_ARG_USAGE,
          buffer=result_array,
          element_type=rt.HalElementType.SINT_32)
      ret_list.push_ref(buffer_view)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi": json.dumps({
            "a": [],
            "r": [["ndarray", "i32", 1, 2]],
        })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker()
    np.testing.assert_array_equal([1, 0], result)

  def testReturnBufferViewNoReflection(self):
    result_array = np.asarray([1, 0], dtype=np.int32)

    def invoke(arg_list, ret_list):
      buffer_view = self.device.allocator.allocate_buffer_copy(
          memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
          allowed_usage=IMPLICIT_BUFFER_ARG_USAGE,
          buffer=result_array,
          element_type=rt.HalElementType.SINT_32)
      ret_list.push_ref(buffer_view)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={})
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker()
    np.testing.assert_array_equal([1, 0], result)

  # TODO: Fill out all return types.
  def testReturnTypeNdArrayBool(self):
    result_array = np.asarray([1, 0], dtype=np.int8)

    def invoke(arg_list, ret_list):
      buffer_view = self.device.allocator.allocate_buffer_copy(
          memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
          allowed_usage=IMPLICIT_BUFFER_ARG_USAGE,
          buffer=result_array,
          element_type=rt.HalElementType.UINT_8)
      ret_list.push_ref(buffer_view)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi": json.dumps({
            "a": [],
            "r": [["ndarray", "i1", 1, 2]],
        })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker()
    # assertEqual on bool arrays is fraught for... reasons.
    np.testing.assert_array_equal([True, False], result)

  def testReturnTypeList(self):
    vm_list = VmVariantList(2)
    vm_list.push_int(1)
    vm_list.push_int(2)

    def invoke(arg_list, ret_list):
      ret_list.push_list(vm_list)

    vm_context = MockVmContext(invoke)
    vm_function = MockVmFunction(reflection={
        "iree.abi":
            json.dumps({
                "a": [],
                "r": [["py_homogeneous_list", "i64"]],
            })
    })
    invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
    result = invoker()
    self.assertEqual("[1, 2]", repr(result))


if __name__ == "__main__":
  unittest.main()