#!/usr/bin/env python3
# Owner(s): ["oncall: mobile"]
# mypy: allow-untyped-defs

import io
import textwrap
from typing import Dict, List, Optional

import torch
import torch.utils.bundled_inputs
from torch.testing._internal.common_utils import run_tests, TestCase


def model_size(sm):
    buffer = io.BytesIO()
    torch.jit.save(sm, buffer)
    return len(buffer.getvalue())


def save_and_load(sm):
    buffer = io.BytesIO()
    torch.jit.save(sm, buffer)
    buffer.seek(0)
    return torch.jit.load(buffer)


class TestBundledInputs(TestCase):
    def test_single_tensors(self):
        class SingleTensorModel(torch.nn.Module):
            def forward(self, arg):
                return arg

        sm = torch.jit.script(SingleTensorModel())
        original_size = model_size(sm)
        get_expr: List[str] = []
        samples = [
            # Tensor with small numel and small storage.
            (torch.tensor([1]),),
            # Tensor with large numel and small storage.
            (torch.tensor([[2, 3, 4]]).expand(1 << 16, -1)[:, ::2],),
            # Tensor with small numel and large storage.
            (torch.tensor(range(1 << 16))[-8:],),
            # Large zero tensor.
            (torch.zeros(1 << 16),),
            # Large channels-last ones tensor.
            (torch.ones(4, 8, 32, 32).contiguous(memory_format=torch.channels_last),),
            # Special encoding of random tensor.
            (torch.utils.bundled_inputs.bundle_randn(1 << 16),),
            # Quantized uniform tensor.
            (torch.quantize_per_tensor(torch.zeros(4, 8, 32, 32), 1, 0, torch.qint8),),
        ]
        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            sm, samples, get_expr
        )
        # print(get_expr[0])
        # print(sm._generate_bundled_inputs.code)

        # Make sure the model only grew a little bit,
        # despite having nominally large bundled inputs.
        augmented_size = model_size(sm)
        self.assertLess(augmented_size, original_size + (1 << 12))

        loaded = save_and_load(sm)
        inflated = loaded.get_all_bundled_inputs()
        self.assertEqual(loaded.get_num_bundled_inputs(), len(samples))
        self.assertEqual(len(inflated), len(samples))
        self.assertTrue(loaded(*inflated[0]) is inflated[0][0])

        for idx, inp in enumerate(inflated):
            self.assertIsInstance(inp, tuple)
            self.assertEqual(len(inp), 1)
            self.assertIsInstance(inp[0], torch.Tensor)
            if idx != 5:
                # Strides might be important for benchmarking.
                self.assertEqual(inp[0].stride(), samples[idx][0].stride())
                self.assertEqual(inp[0], samples[idx][0], exact_dtype=True)

        # This tensor is random, but with 100,000 trials,
        # mean and std had ranges of (-0.0154, 0.0144) and (0.9907, 1.0105).
        self.assertEqual(inflated[5][0].shape, (1 << 16,))
        self.assertEqual(inflated[5][0].mean().item(), 0, atol=0.025, rtol=0)
        self.assertEqual(inflated[5][0].std().item(), 1, atol=0.02, rtol=0)

    def test_large_tensor_with_inflation(self):
        class SingleTensorModel(torch.nn.Module):
            def forward(self, arg):
                return arg

        sm = torch.jit.script(SingleTensorModel())
        sample_tensor = torch.randn(1 << 16)
        # We can store tensors with custom inflation functions regardless
        # of size, even if inflation is just the identity.
        sample = torch.utils.bundled_inputs.bundle_large_tensor(sample_tensor)
        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(sm, [(sample,)])

        loaded = save_and_load(sm)
        inflated = loaded.get_all_bundled_inputs()
        self.assertEqual(len(inflated), 1)

        self.assertEqual(inflated[0][0], sample_tensor)

    def test_rejected_tensors(self):
        def check_tensor(sample):
            # Need to define the class in this scope to get a fresh type for each run.
            class SingleTensorModel(torch.nn.Module):
                def forward(self, arg):
                    return arg

            sm = torch.jit.script(SingleTensorModel())
            with self.assertRaisesRegex(Exception, "Bundled input argument"):
                torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
                    sm, [(sample,)]
                )

        # Plain old big tensor.
        check_tensor(torch.randn(1 << 16))
        # This tensor has two elements, but they're far apart in memory.
        # We currently cannot represent this compactly while preserving
        # the strides.
        small_sparse = torch.randn(2, 1 << 16)[:, 0:1]
        self.assertEqual(small_sparse.numel(), 2)
        check_tensor(small_sparse)

    def test_non_tensors(self):
        class StringAndIntModel(torch.nn.Module):
            def forward(self, fmt: str, num: int):
                return fmt.format(num)

        sm = torch.jit.script(StringAndIntModel())
        samples = [
            ("first {}", 1),
            ("second {}", 2),
        ]
        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(sm, samples)

        loaded = save_and_load(sm)
        inflated = loaded.get_all_bundled_inputs()
        self.assertEqual(inflated, samples)
        self.assertTrue(loaded(*inflated[0]) == "first 1")

    def test_multiple_methods_with_inputs(self):
        class MultipleMethodModel(torch.nn.Module):
            def forward(self, arg):
                return arg

            @torch.jit.export
            def foo(self, arg):
                return arg

        mm = torch.jit.script(MultipleMethodModel())
        samples = [
            # Tensor with small numel and small storage.
            (torch.tensor([1]),),
            # Tensor with large numel and small storage.
            (torch.tensor([[2, 3, 4]]).expand(1 << 16, -1)[:, ::2],),
            # Tensor with small numel and large storage.
            (torch.tensor(range(1 << 16))[-8:],),
            # Large zero tensor.
            (torch.zeros(1 << 16),),
            # Large channels-last ones tensor.
            (torch.ones(4, 8, 32, 32).contiguous(memory_format=torch.channels_last),),
        ]
        info = [
            "Tensor with small numel and small storage.",
            "Tensor with large numel and small storage.",
            "Tensor with small numel and large storage.",
            "Large zero tensor.",
            "Large channels-last ones tensor.",
            "Special encoding of random tensor.",
        ]
        torch.utils.bundled_inputs.augment_many_model_functions_with_bundled_inputs(
            mm,
            inputs={mm.forward: samples, mm.foo: samples},
            info={mm.forward: info, mm.foo: info},
        )
        loaded = save_and_load(mm)
        inflated = loaded.get_all_bundled_inputs()

        # Make sure these functions are all consistent.
        self.assertEqual(inflated, samples)
        self.assertEqual(inflated, loaded.get_all_bundled_inputs_for_forward())
        self.assertEqual(inflated, loaded.get_all_bundled_inputs_for_foo())

        # Check running and size helpers
        self.assertTrue(loaded(*inflated[0]) is inflated[0][0])
        self.assertEqual(loaded.get_num_bundled_inputs(), len(samples))

        # Check helper that work on all functions
        all_info = loaded.get_bundled_inputs_functions_and_info()
        self.assertEqual(set(all_info.keys()), {"forward", "foo"})
        self.assertEqual(
            all_info["forward"]["get_inputs_function_name"],
            ["get_all_bundled_inputs_for_forward"],
        )
        self.assertEqual(
            all_info["foo"]["get_inputs_function_name"],
            ["get_all_bundled_inputs_for_foo"],
        )
        self.assertEqual(all_info["forward"]["info"], info)
        self.assertEqual(all_info["foo"]["info"], info)

        # example of how to turn the 'get_inputs_function_name' into the actual list of bundled inputs
        for func_name in all_info.keys():
            input_func_name = all_info[func_name]["get_inputs_function_name"][0]
            func_to_run = getattr(loaded, input_func_name)
            self.assertEqual(func_to_run(), samples)

    def test_multiple_methods_with_inputs_both_defined_failure(self):
        class MultipleMethodModel(torch.nn.Module):
            def forward(self, arg):
                return arg

            @torch.jit.export
            def foo(self, arg):
                return arg

        samples = [(torch.tensor([1]),)]

        # inputs defined 2 ways so should fail
        with self.assertRaises(Exception):
            mm = torch.jit.script(MultipleMethodModel())
            definition = textwrap.dedent(
                """
                def _generate_bundled_inputs_for_forward(self):
                    return []
                """
            )
            mm.define(definition)
            torch.utils.bundled_inputs.augment_many_model_functions_with_bundled_inputs(
                mm,
                inputs={
                    mm.forward: samples,
                    mm.foo: samples,
                },
            )

    def test_multiple_methods_with_inputs_neither_defined_failure(self):
        class MultipleMethodModel(torch.nn.Module):
            def forward(self, arg):
                return arg

            @torch.jit.export
            def foo(self, arg):
                return arg

        samples = [(torch.tensor([1]),)]

        # inputs not defined so should fail
        with self.assertRaises(Exception):
            mm = torch.jit.script(MultipleMethodModel())
            mm._generate_bundled_inputs_for_forward()
            torch.utils.bundled_inputs.augment_many_model_functions_with_bundled_inputs(
                mm,
                inputs={
                    mm.forward: None,
                    mm.foo: samples,
                },
            )

    def test_bad_inputs(self):
        class SingleTensorModel(torch.nn.Module):
            def forward(self, arg):
                return arg

        # Non list for input list
        with self.assertRaises(TypeError):
            m = torch.jit.script(SingleTensorModel())
            torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
                m,
                inputs="foo",  # type: ignore[arg-type]
            )

        # List of non tuples. Most common error using the api.
        with self.assertRaises(TypeError):
            m = torch.jit.script(SingleTensorModel())
            torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
                m,
                inputs=[torch.ones(1, 2)],  # type: ignore[list-item]
            )

    def test_double_augment_fail(self):
        class SingleTensorModel(torch.nn.Module):
            def forward(self, arg):
                return arg

        m = torch.jit.script(SingleTensorModel())
        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            m, inputs=[(torch.ones(1),)]
        )
        with self.assertRaisesRegex(
            Exception, "Models can only be augmented with bundled inputs once."
        ):
            torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
                m, inputs=[(torch.ones(1),)]
            )

    def test_double_augment_non_mutator(self):
        class SingleTensorModel(torch.nn.Module):
            def forward(self, arg):
                return arg

        m = torch.jit.script(SingleTensorModel())
        bundled_model = torch.utils.bundled_inputs.bundle_inputs(
            m, inputs=[(torch.ones(1),)]
        )
        with self.assertRaises(AttributeError):
            m.get_all_bundled_inputs()
        self.assertEqual(bundled_model.get_all_bundled_inputs(), [(torch.ones(1),)])
        self.assertEqual(bundled_model.forward(torch.ones(1)), torch.ones(1))

    def test_double_augment_success(self):
        class SingleTensorModel(torch.nn.Module):
            def forward(self, arg):
                return arg

        m = torch.jit.script(SingleTensorModel())
        bundled_model = torch.utils.bundled_inputs.bundle_inputs(
            m, inputs={m.forward: [(torch.ones(1),)]}
        )
        self.assertEqual(bundled_model.get_all_bundled_inputs(), [(torch.ones(1),)])

        bundled_model2 = torch.utils.bundled_inputs.bundle_inputs(
            bundled_model, inputs=[(torch.ones(2),)]
        )
        self.assertEqual(bundled_model2.get_all_bundled_inputs(), [(torch.ones(2),)])

    def test_dict_args(self):
        class MyModel(torch.nn.Module):
            def forward(
                self,
                arg1: Optional[Dict[str, torch.Tensor]],
                arg2: Optional[List[torch.Tensor]],
                arg3: torch.Tensor,
            ):
                if arg1 is None:
                    return arg3
                elif arg2 is None:
                    return arg1["a"] + arg1["b"]
                else:
                    return arg1["a"] + arg1["b"] + arg2[0]

        small_sample = dict(
            a=torch.zeros([10, 20]),
            b=torch.zeros([1, 1]),
            c=torch.zeros([10, 20]),
        )
        small_list = [torch.zeros([10, 20])]

        big_sample = dict(
            a=torch.zeros([1 << 5, 1 << 8, 1 << 10]),
            b=torch.zeros([1 << 5, 1 << 8, 1 << 10]),
            c=torch.zeros([1 << 5, 1 << 8, 1 << 10]),
        )
        big_list = [torch.zeros([1 << 5, 1 << 8, 1 << 10])]

        def condensed(t):
            ret = torch.empty_like(t).flatten()[0].clone().expand(t.shape)
            assert ret.storage().size() == 1
            # ret.storage()[0] = 0
            return ret

        def bundle_optional_dict_of_randn(template):
            return torch.utils.bundled_inputs.InflatableArg(
                value=(
                    None
                    if template is None
                    else {k: condensed(v) for (k, v) in template.items()}
                ),
                fmt="{}",
                fmt_fn="""
                def {}(self, value: Optional[Dict[str, Tensor]]):
                    if value is None:
                        return None
                    output = {{}}
                    for k, v in value.items():
                        output[k] = torch.randn_like(v)
                    return output
                """,
            )

        def bundle_optional_list_of_randn(template):
            return torch.utils.bundled_inputs.InflatableArg(
                value=(None if template is None else [condensed(v) for v in template]),
                fmt="{}",
                fmt_fn="""
                def {}(self, value: Optional[List[Tensor]]):
                    if value is None:
                        return None
                    output = []
                    for v in value:
                        output.append(torch.randn_like(v))
                    return output
                """,
            )

        out: List[str] = []
        sm = torch.jit.script(MyModel())
        original_size = model_size(sm)
        small_inputs = (
            bundle_optional_dict_of_randn(small_sample),
            bundle_optional_list_of_randn(small_list),
            torch.zeros([3, 4]),
        )
        big_inputs = (
            bundle_optional_dict_of_randn(big_sample),
            bundle_optional_list_of_randn(big_list),
            torch.zeros([1 << 5, 1 << 8, 1 << 10]),
        )

        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            sm,
            [big_inputs, small_inputs],
            _receive_inflate_expr=out,
        )
        augmented_size = model_size(sm)
        # assert the size has not increased more than 8KB
        self.assertLess(augmented_size, original_size + (1 << 13))

        loaded = save_and_load(sm)
        inflated = loaded.get_all_bundled_inputs()
        self.assertEqual(len(inflated[0]), len(small_inputs))

        methods, _ = (
            torch.utils.bundled_inputs._get_bundled_inputs_attributes_and_methods(
                loaded
            )
        )

        # One Function (forward)
        # two bundled inputs (big_inputs and small_inputs)
        # two args which have InflatableArg with fmt_fn
        # 1 * 2 * 2 = 4
        self.assertEqual(
            sum(method.startswith("_inflate_helper") for method in methods), 4
        )


if __name__ == "__main__":
    run_tests()