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

import os
import sys
import torch
from torch.testing._internal.jit_utils import JitTestCase, make_global
from torch.jit._monkeytype_config import _IS_MONKEYTYPE_INSTALLED
from typing import List, Dict, Tuple, Any, Optional, NamedTuple  # noqa: F401
from torch.testing._internal.common_utils import NoTest

# 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):
                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), ))