import os
import sys
import unittest

import torch
import torch._C
from pathlib import Path
from test_nnapi import TestNNAPI
from torch.testing._internal.common_utils import TEST_WITH_ASAN

# Make the helper files in test/ importable
pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(pytorch_test_dir)

if __name__ == "__main__":
    raise RuntimeError(
        "This test file is not meant to be run directly, use:\n\n"
        "\tpython test/test_jit.py TESTNAME\n\n"
        "instead."
    )

"""
Unit Tests for Nnapi backend with delegate
Inherits most tests from TestNNAPI, which loads Android NNAPI models
without the delegate API.
"""
# First skip is needed for IS_WINDOWS or IS_MACOS to skip the tests.
# Second skip is because ASAN is currently causing an error.
# It is still unclear how to resolve this. T95764916
torch_root = Path(__file__).resolve().parent.parent.parent
lib_path = torch_root / 'build' / 'lib' / 'libnnapi_backend.so'
@unittest.skipIf(not os.path.exists(lib_path),
                 "Skipping the test as libnnapi_backend.so was not found")
@unittest.skipIf(TEST_WITH_ASAN, "Unresolved bug with ASAN")
class TestNnapiBackend(TestNNAPI):
    def setUp(self):
        super().setUp()

        # Save default dtype
        module = torch.nn.PReLU()
        self.default_dtype = module.weight.dtype
        # Change dtype to float32 (since a different unit test changed dtype to float64,
        # which is not supported by the Android NNAPI delegate)
        # Float32 should typically be the default in other files.
        torch.set_default_dtype(torch.float32)

        # Load nnapi delegate library
        torch.ops.load_library(str(lib_path))

    # Override
    def call_lowering_to_nnapi(self, traced_module, args):
        compile_spec = {"forward": {"inputs": args}}
        return torch._C._jit_to_backend("nnapi", traced_module, compile_spec)

    def test_tensor_input(self):
        # Lower a simple module
        args = torch.tensor([[1.0, -1.0, 2.0, -2.0]]).unsqueeze(-1).unsqueeze(-1)
        module = torch.nn.PReLU()
        traced = torch.jit.trace(module, args)

        # Argument input is a single Tensor
        self.call_lowering_to_nnapi(traced, args)
        # Argument input is a Tensor in a list
        self.call_lowering_to_nnapi(traced, [args])

    # Test exceptions for incorrect compile specs
    def test_compile_spec_santiy(self):
        args = torch.tensor([[1.0, -1.0, 2.0, -2.0]]).unsqueeze(-1).unsqueeze(-1)
        module = torch.nn.PReLU()
        traced = torch.jit.trace(module, args)

        errorMsgTail = r"""
method_compile_spec should contain a Tensor or Tensor List which bundles input parameters: shape, dtype, quantization, and dimorder.
For input shapes, use 0 for run/load time flexible input.
method_compile_spec must use the following format:
{"forward": {"inputs": at::Tensor}} OR {"forward": {"inputs": c10::List<at::Tensor>}}"""

        # No forward key
        compile_spec = {"backward": {"inputs": args}}
        with self.assertRaisesRegex(RuntimeError, "method_compile_spec does not contain the \"forward\" key." + errorMsgTail):
            torch._C._jit_to_backend("nnapi", traced, compile_spec)

        # No dictionary under the forward key
        compile_spec = {"forward": 1}
        with self.assertRaisesRegex(RuntimeError,
                                    "method_compile_spec does not contain a dictionary with an \"inputs\" key, "
                                    "under it's \"forward\" key."
                                    + errorMsgTail):
            torch._C._jit_to_backend("nnapi", traced, compile_spec)

        # No inputs key (in the dictionary under the forward key)
        compile_spec = {"forward": {"not inputs": args}}
        with self.assertRaisesRegex(RuntimeError,
                                    "method_compile_spec does not contain a dictionary with an \"inputs\" key, "
                                    "under it's \"forward\" key."
                                    + errorMsgTail):
            torch._C._jit_to_backend("nnapi", traced, compile_spec)

        # No Tensor or TensorList under the inputs key
        compile_spec = {"forward": {"inputs": 1}}
        with self.assertRaisesRegex(RuntimeError,
                                    "method_compile_spec does not contain either a Tensor or TensorList, under it's \"inputs\" key."
                                    + errorMsgTail):
            torch._C._jit_to_backend("nnapi", traced, compile_spec)
        compile_spec = {"forward": {"inputs": [1]}}
        with self.assertRaisesRegex(RuntimeError,
                                    "method_compile_spec does not contain either a Tensor or TensorList, under it's \"inputs\" key."
                                    + errorMsgTail):
            torch._C._jit_to_backend("nnapi", traced, compile_spec)

    def tearDown(self):
        # Change dtype back to default (Otherwise, other unit tests will complain)
        torch.set_default_dtype(self.default_dtype)

[((13, 0, 13, 33), 'sys.path.append', 'sys.path.append', ({(13, 16, 13, 32): 'pytorch_test_dir'}, {}), '(pytorch_test_dir)', False, 'import sys\n'), ((34, 1, 34, 60), 'unittest.skipIf', 'unittest.skipIf', ({(34, 17, 34, 31): 'TEST_WITH_ASAN', (34, 33, 34, 59): '"""Unresolved bug with ASAN"""'}, {}), "(TEST_WITH_ASAN, 'Unresolved bug with ASAN')", False, 'import unittest\n'), ((12, 51, 12, 77), 'os.path.realpath', 'os.path.realpath', ({(12, 68, 12, 76): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((40, 17, 40, 33), 'torch.nn.PReLU', 'torch.nn.PReLU', ({}, {}), '()', False, 'import torch\n'), ((45, 8, 45, 46), 'torch.set_default_dtype', 'torch.set_default_dtype', ({(45, 32, 45, 45): 'torch.float32'}, {}), '(torch.float32)', False, 'import torch\n'), ((53, 15, 53, 77), 'torch._C._jit_to_backend', 'torch._C._jit_to_backend', ({(53, 40, 53, 47): '"""nnapi"""', (53, 49, 53, 62): 'traced_module', (53, 64, 53, 76): 'compile_spec'}, {}), "('nnapi', traced_module, compile_spec)", False, 'import torch\n'), ((58, 17, 58, 33), 'torch.nn.PReLU', 'torch.nn.PReLU', ({}, {}), '()', False, 'import torch\n'), ((59, 17, 59, 46), 'torch.jit.trace', 'torch.jit.trace', ({(59, 33, 59, 39): 'module', (59, 41, 59, 45): 'args'}, {}), '(module, args)', False, 'import torch\n'), ((69, 17, 69, 33), 'torch.nn.PReLU', 'torch.nn.PReLU', ({}, {}), '()', False, 'import torch\n'), ((70, 17, 70, 46), 'torch.jit.trace', 'torch.jit.trace', ({(70, 33, 70, 39): 'module', (70, 41, 70, 45): 'args'}, {}), '(module, args)', False, 'import torch\n'), ((113, 8, 113, 51), 'torch.set_default_dtype', 'torch.set_default_dtype', ({(113, 32, 113, 50): 'self.default_dtype'}, {}), '(self.default_dtype)', False, 'import torch\n'), ((32, 21, 32, 45), 'os.path.exists', 'os.path.exists', ({(32, 36, 32, 44): 'lib_path'}, {}), '(lib_path)', False, 'import os\n'), ((81, 12, 81, 67), 'torch._C._jit_to_backend', 'torch._C._jit_to_backend', ({(81, 37, 81, 44): '"""nnapi"""', (81, 46, 81, 52): 'traced', (81, 54, 81, 66): 'compile_spec'}, {}), "('nnapi', traced, compile_spec)", False, 'import torch\n'), ((89, 12, 89, 67), 'torch._C._jit_to_backend', 'torch._C._jit_to_backend', ({(89, 37, 89, 44): '"""nnapi"""', (89, 46, 89, 52): 'traced', (89, 54, 89, 66): 'compile_spec'}, {}), "('nnapi', traced, compile_spec)", False, 'import torch\n'), ((97, 12, 97, 67), 'torch._C._jit_to_backend', 'torch._C._jit_to_backend', ({(97, 37, 97, 44): '"""nnapi"""', (97, 46, 97, 52): 'traced', (97, 54, 97, 66): 'compile_spec'}, {}), "('nnapi', traced, compile_spec)", False, 'import torch\n'), ((104, 12, 104, 67), 'torch._C._jit_to_backend', 'torch._C._jit_to_backend', ({(104, 37, 104, 44): '"""nnapi"""', (104, 46, 104, 52): 'traced', (104, 54, 104, 66): 'compile_spec'}, {}), "('nnapi', traced, compile_spec)", False, 'import torch\n'), ((109, 12, 109, 67), 'torch._C._jit_to_backend', 'torch._C._jit_to_backend', ({(109, 37, 109, 44): '"""nnapi"""', (109, 46, 109, 52): 'traced', (109, 54, 109, 66): 'compile_spec'}, {}), "('nnapi', traced, compile_spec)", False, 'import torch\n'), ((30, 13, 30, 27), 'pathlib.Path', 'Path', ({(30, 18, 30, 26): '__file__'}, {}), '(__file__)', False, 'from pathlib import Path\n'), ((57, 15, 57, 53), 'torch.tensor', 'torch.tensor', ({(57, 28, 57, 52): '[[1.0, -1.0, 2.0, -2.0]]'}, {}), '([[1.0, -1.0, 2.0, -2.0]])', False, 'import torch\n'), ((68, 15, 68, 53), 'torch.tensor', 'torch.tensor', ({(68, 28, 68, 52): '[[1.0, -1.0, 2.0, -2.0]]'}, {}), '([[1.0, -1.0, 2.0, -2.0]])', False, 'import torch\n')]