# Owner(s): ["oncall: jit"]

import os
import sys
from typing import Any, Dict, List, NamedTuple, Optional, Tuple  # noqa: F401

import torch
from torch.jit._monkeytype_config import _IS_MONKEYTYPE_INSTALLED
from torch.testing._internal.common_utils import NoTest
from torch.testing._internal.jit_utils import JitTestCase, make_global


# 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 not _IS_MONKEYTYPE_INSTALLED:
    print(
        "monkeytype is not installed. Skipping tests for Profile-Directed Typing",
        file=sys.stderr,
    )
    JitTestCase = NoTest  # type: ignore[misc, assignment] # noqa: F811

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."
    )


class TestPDT(JitTestCase):
    """
    A suite of tests for profile directed typing in TorchScript.
    """

    def test_nn_module(self):
        class TestPDTModel(torch.nn.Module):
            def forward(self, x) -> Any:
                if isinstance(x, int):
                    return x + 1
                elif isinstance(x, float):
                    return x - 1
                else:
                    return x

        make_global(TestPDTModel)
        pdt_model = TestPDTModel()
        inp: List[Tuple[Any, ...]] = [
            (20,),
            (2.7,),
            (False,),
        ]
        scripted_pdt_model = torch.jit.script(
            pdt_model, example_inputs={pdt_model: inp}
        )
        self.assertEqual(scripted_pdt_model(50), pdt_model(50))
        self.assertEqual(scripted_pdt_model(1.8), pdt_model(1.8))
        self.assertTrue(scripted_pdt_model(True), pdt_model(True))

    def test_nested_nn_module_class(self):
        class NestedPDTInner(torch.nn.Module):
            def forward(self, x):
                if isinstance(x, int):
                    return x * 10
                return x

        class NestedModulePDTWrapper(torch.nn.Module):
            def __init__(self, inner):
                super().__init__()
                self.inner = inner

            def forward(self, x):
                return self.inner(x)

        make_global(NestedPDTInner, NestedModulePDTWrapper)
        inner_pdt_model = NestedPDTInner()
        wrapped_pdt_model = NestedModulePDTWrapper(inner_pdt_model)
        inp: List[Tuple[Any, ...]] = [(20,), (False,)]
        scripted_pdt_model = torch.jit.script(
            wrapped_pdt_model, example_inputs={wrapped_pdt_model: inp}
        )
        self.assertEqual(scripted_pdt_model(30), wrapped_pdt_model(30))
        self.assertEqual(scripted_pdt_model(1.9), wrapped_pdt_model(1.9))
        self.assertTrue(scripted_pdt_model(True), wrapped_pdt_model(True))

    def test_nested_nn_module_class_with_args(self):
        class NestedModulePDTInner(torch.nn.Module):
            def forward(self, x, y):
                if isinstance(x, int):
                    return x * 10 + y
                return x

        class NestedModulePDTOuter(torch.nn.Module):
            def __init__(self, inner):
                super().__init__()
                self.inner = inner

            def forward(self, x):
                return self.inner(x, 20)

        make_global(NestedModulePDTInner, NestedModulePDTOuter)
        inner_pdt_model = NestedModulePDTInner()
        outer_pdt_model = NestedModulePDTOuter(inner_pdt_model)
        inner_input: List[Tuple[Any, ...]] = [
            (10, 10),
            (1.9, 20),
        ]
        outer_input: List[Tuple[Any, ...]] = [(20,), (False,)]
        scripted_pdt_model = torch.jit.script(
            outer_pdt_model,
            example_inputs={
                inner_pdt_model: inner_input,
                outer_pdt_model: outer_input,
            },
        )
        self.assertEqual(scripted_pdt_model(30), outer_pdt_model(30))
        self.assertEqual(scripted_pdt_model(1.9), outer_pdt_model(1.9))
        self.assertTrue(scripted_pdt_model(True), outer_pdt_model(True))

    def test_nested_function_in_forward(self):
        class NestedFunctionInForward(torch.nn.Module):
            def forward(self, x):
                return self.fun(x) + 10

            def fun(self, x):
                if isinstance(x, bool):
                    return 0
                elif isinstance(x, int):
                    return x + 1
                return 0

        make_global(NestedFunctionInForward)
        pdt_model = NestedFunctionInForward()
        inp: List[Tuple[Any, ...]] = [(-1,), (False,)]
        scripted_pdt_model = torch.jit.script(
            pdt_model, example_inputs={pdt_model: inp}
        )
        self.assertEqual(scripted_pdt_model(30), pdt_model(30))
        self.assertEqual(scripted_pdt_model(True), pdt_model(True))

    def test_nn_module_with_export_function(self):
        class TestModelWithExport(torch.nn.Module):
            @torch.jit.export
            def fn(self, x, y) -> Any:
                assert not (isinstance(x, bool) and isinstance(y, bool))
                if isinstance(x, int) and isinstance(y, int):
                    return x + y
                elif isinstance(x, float) and isinstance(y, float):
                    return x - y
                else:
                    return -1

        make_global(TestModelWithExport)
        pdt_model = TestModelWithExport()
        inp: List[Tuple[Any, ...]] = [
            (
                20,
                10,
            ),
            (
                2.7,
                8.9,
            ),
        ]
        scripted_pdt_model = torch.jit.script(
            pdt_model, example_inputs={pdt_model.fn: inp}
        )
        self.assertEqual(scripted_pdt_model.fn(10, 90), pdt_model.fn(10, 90))
        self.assertEqual(scripted_pdt_model.fn(1.8, 2.2), pdt_model.fn(1.8, 2.2))
        self.assertTrue(
            scripted_pdt_model.fn(torch.ones(1), 2), pdt_model.fn(torch.ones(1), 2)
        )

    def test_class_methods(self):
        class PDTModel:
            def test_sum(self, a):
                return sum(a)

        make_global(PDTModel)
        pdt_model = PDTModel()
        inp: List[Tuple[Any, ...]] = [
            (
                [
                    10,
                    20,
                ],
            ),
        ]
        scripted_pdt_model = torch.jit.script(
            PDTModel, example_inputs={pdt_model.test_sum: inp}
        )
        script_model = scripted_pdt_model()
        self.assertEqual(
            script_model.test_sum(
                [
                    10,
                    20,
                    30,
                ],
            ),
            pdt_model.test_sum(
                [
                    10,
                    20,
                    30,
                ],
            ),
        )

    def test_class_with_multiple_methods(self):
        class PDTModelWithManyMethods:
            def test_list_to_dict(self, a):
                new_dictionary: Dict[float, bool] = {}
                for element in a:
                    new_dictionary[element] = True
                return new_dictionary

            def test_substring(self, a, b):
                return b in a

        make_global(PDTModelWithManyMethods)
        pdt_model = PDTModelWithManyMethods()
        list_inp: List[Tuple[Any, ...]] = [
            (
                [
                    1.2,
                    2.3,
                ],
            ),
        ]
        str_inp: List[Tuple[Any, ...]] = [
            (
                "abc",
                "b",
            ),
        ]
        scripted_pdt_model = torch.jit.script(
            PDTModelWithManyMethods,
            example_inputs={
                pdt_model.test_list_to_dict: list_inp,
                pdt_model.test_substring: str_inp,
            },
        )
        script_model = scripted_pdt_model()
        self.assertEqual(
            script_model.test_list_to_dict(
                [
                    1.1,
                    2.2,
                    3.3,
                ],
            ),
            pdt_model.test_list_to_dict(
                [
                    1.1,
                    2.2,
                    3.3,
                ],
            ),
        )
        self.assertEqual(
            script_model.test_substring(
                "helloworld",
                "world",
            ),
            pdt_model.test_substring(
                "helloworld",
                "world",
            ),
        )
        self.assertEqual(
            script_model.test_substring(
                "helloworld",
                "def",
            ),
            pdt_model.test_substring(
                "helloworld",
                "def",
            ),
        )

    def test_multiple_class_with_same_method(self):
        class PDTModelOne:
            def test_find(self, a, b):
                return b in a.keys()

        class PDTModelTwo:
            def test_find(self, a, b):
                return b in a

        make_global(PDTModelOne, PDTModelTwo)
        pdt_model_one = PDTModelOne()
        pdt_model_two = PDTModelTwo()
        dict_inp: List[Tuple[Any, ...]] = [
            (
                {
                    1.2: True,
                    2.3: False,
                },
                1.2,
            ),
        ]
        list_inp: List[Tuple[Any, ...]] = [
            (
                [
                    "abc",
                    "b",
                ],
                "c",
            ),
        ]
        scripted_pdt_model_one = torch.jit.script(
            PDTModelOne, example_inputs={pdt_model_one.test_find: dict_inp}
        )
        scripted_pdt_model_two = torch.jit.script(
            PDTModelTwo, example_inputs={pdt_model_two.test_find: list_inp}
        )

        script_model_one, script_model_two = (
            scripted_pdt_model_one(),
            scripted_pdt_model_two(),
        )
        self.assertEqual(
            script_model_one.test_find(
                {
                    1.1: True,
                    2.2: True,
                    3.3: False,
                },
                4.4,
            ),
            pdt_model_one.test_find(
                {
                    1.1: True,
                    2.2: True,
                    3.3: False,
                },
                4.4,
            ),
        )
        self.assertEqual(
            script_model_two.test_find(
                [
                    "hello",
                    "world",
                ],
                "world",
            ),
            pdt_model_two.test_find(
                [
                    "hello",
                    "world",
                ],
                "world",
            ),
        )

    def test_pdt(self):
        def test_sum(a, b):
            return a + b

        make_global(test_sum)
        scripted_fn_add = torch.jit.script(test_sum, example_inputs=[(3, 4)])
        self.assertEqual(scripted_fn_add(10, 2), test_sum(10, 2))

        def test_sub(a, b):
            return a - b

        make_global(test_sub)
        scripted_fn_sub = torch.jit.script(test_sub, example_inputs=[(3.9, 4.10)])
        self.assertEqual(scripted_fn_sub(6.5, 2.9), test_sub(6.5, 2.9))

        def test_mul(a, b):
            return a * b

        make_global(test_mul)
        scripted_fn_mul = torch.jit.script(test_mul, example_inputs=[(-10, 9)])
        self.assertEqual(scripted_fn_mul(-1, 3), test_mul(-1, 3))

        def test_args_complex(real, img):
            return torch.complex(real, img)

        make_global(test_args_complex)
        scripted_fn_complex = torch.jit.script(
            test_args_complex, example_inputs=[(torch.rand(3, 4), torch.rand(3, 4))]
        )
        arg1, arg2 = torch.rand(3, 4), torch.rand(3, 4)
        self.assertEqual(scripted_fn_complex(arg1, arg2), test_args_complex(arg1, arg2))

        def test_bool(a):
            if a:
                return -1
            else:
                return 0

        make_global(test_bool)
        scripted_fn_bool = torch.jit.script(test_bool, example_inputs=[(True,)])
        self.assertEqual(scripted_fn_bool(True), test_bool(True))

        def test_str(a):
            if a == "":
                return False
            else:
                return True

        make_global(test_str)
        scripted_fn_str = torch.jit.script(test_str, example_inputs=[("",)])
        self.assertEqual(scripted_fn_str("abc"), test_str("abc"))

    def test_pdt_list_and_tuple(self):
        def test_list_and_tuple(a):
            return sum(a)

        make_global(test_list_and_tuple)

        scripted_fn_float_list_input = torch.jit.script(
            test_list_and_tuple, example_inputs=[([4.9, 8.9],)]
        )
        self.assertEqual(
            scripted_fn_float_list_input([11.9, 7.6]), test_list_and_tuple([11.9, 7.6])
        )

        scripted_fn_bool_list_input = torch.jit.script(
            test_list_and_tuple, example_inputs=[([True, False, True],)]
        )
        self.assertEqual(
            scripted_fn_bool_list_input([True, True, True]),
            test_list_and_tuple([True, True, True]),
        )

        scripted_fn_int_list_input = torch.jit.script(
            test_list_and_tuple, example_inputs=[([3, 4, 5],)]
        )
        self.assertEqual(
            scripted_fn_int_list_input([1, 2, 3]), test_list_and_tuple([1, 2, 3])
        )

        scripted_fn_float_tuple_input = torch.jit.script(
            test_list_and_tuple, example_inputs=[((4.9, 8.9),)]
        )
        self.assertEqual(
            scripted_fn_float_tuple_input((11.9, 7.6)), test_list_and_tuple((11.9, 7.6))
        )

        scripted_fn_bool_tuple_input = torch.jit.script(
            test_list_and_tuple, example_inputs=[((True, False, True),)]
        )
        self.assertEqual(
            scripted_fn_bool_tuple_input((True, True, True)),
            test_list_and_tuple((True, True, True)),
        )

        scripted_fn_int_tuple_input = torch.jit.script(
            test_list_and_tuple, example_inputs=[((3, 4, 5),)]
        )
        self.assertEqual(
            scripted_fn_int_tuple_input((1, 2, 3)), test_list_and_tuple((1, 2, 3))
        )

    def test_nested_list_and_tuple(self):
        def test_nested_list(inp):
            return [sum(v) for v in inp]

        def test_nested_tuple(inp):
            ans = 0.0
            for tup in inp:
                for val in tup:
                    if val > 0:
                        ans *= val
            return ans

        make_global(test_nested_list, test_nested_tuple)

        list_inp = [
            [
                1,
                2,
                3,
            ],
            [
                5,
                6,
                7,
            ],
        ]
        scripted_fn = torch.jit.script(
            test_nested_list,
            example_inputs=[
                (list_inp,),
            ],
        )
        inp = [
            [
                0,
                4,
                7,
            ],
            [
                8,
                11,
            ],
            [
                6,
                -1,
                -20,
            ],
        ]
        self.assertEqual(
            scripted_fn(
                inp,
            ),
            test_nested_list(
                inp,
            ),
        )

        list_inp = (
            [
                1,
                2,
                3,
            ],
            [
                5,
                6,
                7,
            ],
        )
        scripted_fn = torch.jit.script(
            test_nested_list,
            example_inputs=[
                (list_inp,),
            ],
        )
        inp = (
            [
                0,
                4,
                7,
            ],
            [
                8,
                11,
            ],
            [
                6,
                -1,
                -20,
            ],
        )
        self.assertEqual(
            scripted_fn(
                inp,
            ),
            test_nested_list(
                inp,
            ),
        )

        tup_inp = [
            (
                1.0,
                2.6,
                3.7,
            ),
            (
                5.7,
                6.1,
                1.7,
            ),
        ]
        scripted_fn = torch.jit.script(
            test_nested_tuple,
            example_inputs=[
                (tup_inp,),
            ],
        )
        inp = [
            (
                1.0,
                4.1,
                7.4,
            ),
            (
                4.8,
                1.1,
                -1.2,
            ),
            (
                6.3,
                -1.3,
                -2.0,
            ),
        ]
        self.assertEqual(
            scripted_fn(
                inp,
            ),
            test_nested_tuple(
                inp,
            ),
        )

        tup_inp = (
            (
                True,
                False,
                True,
            ),
            (
                False,
                False,
                False,
            ),
        )
        scripted_fn = torch.jit.script(
            test_nested_tuple,
            example_inputs=[
                (tup_inp,),
            ],
        )
        inp = (
            (
                True,
                True,
                True,
            ),
            (
                False,
                False,
                True,
            ),
        )
        self.assertEqual(
            scripted_fn(
                inp,
            ),
            test_nested_tuple(
                inp,
            ),
        )

    def test_pdt_dict(self):
        def test_dict(a):
            return a["foo"]

        def test_dict_int_list(a):
            return a[1]

        make_global(test_dict, test_dict_int_list)

        str_bool_inp = {"foo": True, "bar": False}
        scripted_fn = torch.jit.script(test_dict, example_inputs=[(str_bool_inp,)])
        self.assertEqual(
            scripted_fn(
                {"foo": False, "bar": True},
            ),
            test_dict(
                {"foo": False, "bar": True},
            ),
        )

        str_list_inp = {0: [True, False], 1: [False, True]}
        scripted_fn = torch.jit.script(
            test_dict_int_list, example_inputs=[(str_list_inp,)]
        )
        self.assertEqual(
            scripted_fn(
                {0: [False, False], 1: [True, True]},
            ),
            test_dict_int_list(
                {0: [False, False], 1: [True, True]},
            ),
        )

    def test_any(self):
        def test_multiple_types(a):
            assert not isinstance(a, bool)
            return a

        def test_multiple_type_refinement(a):
            if isinstance(a, bool):
                return 1
            elif isinstance(a, int):
                return 1 + a
            elif isinstance(a, float):
                return 1 + int(a)
            else:
                return -1

        make_global(test_multiple_types, test_multiple_type_refinement)

        scripted_fn = torch.jit.script(
            test_multiple_types, example_inputs=[(1,), ("abc",), (8.9,), ([3, 4, 5],)]
        )
        self.assertEqual(scripted_fn(10), test_multiple_types(10))
        self.assertEqual(scripted_fn("def"), test_multiple_types("def"))
        self.assertEqual(scripted_fn(7.89999), test_multiple_types(7.89999))
        self.assertEqual(scripted_fn([10, 11, 14]), test_multiple_types([10, 11, 14]))

        scripted_fn = torch.jit.script(
            test_multiple_type_refinement,
            example_inputs=[
                (1,),
                ("abc",),
                (8.9,),
                ([3, 4, 5],),
                (True,),
                ({"a": True},),
            ],
        )
        self.assertEqual(scripted_fn(10), test_multiple_type_refinement(10))
        self.assertEqual(scripted_fn("def"), test_multiple_type_refinement("def"))
        self.assertEqual(scripted_fn(7.89999), test_multiple_type_refinement(7.89999))
        self.assertEqual(
            scripted_fn([10, 11, 14]), test_multiple_type_refinement([10, 11, 14])
        )
        self.assertEqual(scripted_fn(False), test_multiple_type_refinement(False))
        self.assertEqual(
            scripted_fn({"abc": True, "def": False}),
            test_multiple_type_refinement({"abc": True, "def": False}),
        )

    def test_class_as_profiled_types(self):
        class UserDefinedClass:
            def fn(self, b) -> Any:
                assert b is not None
                if isinstance(b, int):
                    return b if b > 0 else -1
                elif isinstance(b, float):
                    return b if b > 0.0 else -1.0
                return 0

        def test_model(a, m):
            assert not isinstance(a, bool)
            return m.fn(a)

        make_global(UserDefinedClass, test_model)

        user_class = UserDefinedClass()
        scripted_fn = torch.jit.script(
            test_model,
            example_inputs=[
                (
                    10,
                    user_class,
                ),
                (
                    10.9,
                    user_class,
                ),
            ],
        )
        self.assertEqual(
            scripted_fn(
                100,
                user_class,
            ),
            test_model(100, user_class),
        )
        self.assertEqual(
            scripted_fn(
                1.9,
                user_class,
            ),
            test_model(1.9, user_class),
        )

    def test_class_with_args_as_profiled_types(self):
        class ClassWithArgs:
            def __init__(self, a: bool):
                self.a = a

            def fn(self, b):
                if self.a:
                    return b
                else:
                    return -1

        def test_model_with_args(a, m):
            assert not isinstance(a, bool)
            return m.fn(a)

        make_global(ClassWithArgs, test_model_with_args)

        user_class = ClassWithArgs(False)
        scripted_fn = torch.jit.script(
            test_model_with_args,
            example_inputs=[
                (
                    10,
                    user_class,
                ),
                (
                    10.9,
                    user_class,
                ),
            ],
        )
        self.assertEqual(
            scripted_fn(
                100,
                ClassWithArgs(True),
            ),
            test_model_with_args(100, ClassWithArgs(True)),
        )

    def test_nn_parameter_as_arg(self):
        class TestNNParameter(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.inp = torch.nn.Parameter(torch.ones(2, 3))

            def add_nn_parameter_with_int(self, x, y):
                return torch.add(x, y)

            def forward(self, y):
                return self.add_nn_parameter_with_int(self.inp, y)

        make_global(TestNNParameter)
        pdt_model = TestNNParameter()
        scripted_fn = torch.jit.script(
            pdt_model,
            example_inputs={
                pdt_model: [
                    (10,),
                ],
            },
        )
        self.assertEqual(scripted_fn(20), pdt_model(20))

    def test_fx_tracing_with_typing(self):
        class FXModelOutput(NamedTuple):
            result: List[int]

        class FXModel(torch.nn.Module):
            def forward(self, a) -> FXModelOutput:
                result = FXModelOutput(result=a)
                return result

        make_global(FXModel, FXModelOutput)
        pdt_model = FXModel()
        scripted_fn = torch.jit.script(
            pdt_model,
            example_inputs={
                pdt_model: [
                    (
                        [
                            10,
                            20,
                        ],
                    ),
                ],
            },
        )
        self.assertEqual(scripted_fn([20]), pdt_model([20]))

    def test_nonetype_as_optional_of_type(self):
        def test_none(a) -> Any:
            if a is None:
                return 0
            else:
                return a + torch.ones(1)

        make_global(test_none)

        scripted_fn = torch.jit.script(test_none, example_inputs=[(None,), (10.6,)])
        self.assertEqual(
            scripted_fn(
                30.9,
            ),
            test_none(
                30.9,
            ),
        )

        scripted_fn = torch.jit.script(test_none, example_inputs=[(None,), (10,)])
        self.assertEqual(
            scripted_fn(
                2,
            ),
            test_none(
                2,
            ),
        )

        scripted_fn = torch.jit.script(
            test_none, example_inputs=[(None,), (torch.Tensor(1),)]
        )
        self.assertEqual(
            scripted_fn(
                torch.ones(1),
            ),
            test_none(
                torch.ones(1),
            ),
        )