import json
from typing import Iterable

import torch

from elasticai.creator.precomputation.precomputation import (
    JSONEncoder,
    Precomputation,
    precomputable,
)
from elasticai.creator.tags_utils import Module, get_tags
from elasticai.creator.tests.tensor_test_case import TensorTestCase


class DummyModule:
    def __init__(self):
        self.call = lambda x: x

    @property
    def training(self) -> bool:
        return False

    # noinspection PyMethodMayBeStatic
    def extra_repr(self) -> str:
        return ""

    # noinspection PyMethodMayBeStatic
    def named_children(self) -> Iterable[tuple[str, "Module"]]:
        yield from ()

    def __call__(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        return self.call(x)


class PrecomputationTest(TensorTestCase):
    def test_precompute(self):
        module = DummyModule()

        def call(x: torch.Tensor) -> torch.Tensor:
            return x * 2

        module.call = call
        precompute = Precomputation(
            module=module, input_domain=torch.tensor([[1], [1]])
        )
        precompute()
        self.assertTensorEquals(torch.tensor([[2], [2]]), precompute.output)

    def test_precomputation_is_json_encodable(self):
        module = DummyModule()
        precompute = Precomputation(module=module, input_domain=torch.tensor([[]]))
        precompute()
        actual = json.dumps(precompute, cls=JSONEncoder)
        expected = """{"description": [], "shape": [0], "x": [[]], "y": [[]]}"""
        self.assertEqual(expected, actual)

    def test_precomputation_is_json_encodable_for_shape1x1(self):
        module = DummyModule()
        precompute = Precomputation(module=module, input_domain=torch.tensor([[[1]]]))
        precompute()
        actual = json.dumps(precompute, cls=JSONEncoder)
        expected = (
            """{"description": [], "shape": [1, 1], "x": [[[1]]], "y": [[[1]]]}"""
        )
        self.assertEqual(expected, actual)

    def test_precomputation_is_json_encoding_correctly_writes_in_and_output(self):
        module = DummyModule()

        def call(x: torch.Tensor) -> torch.Tensor:
            return torch.flatten(x * 3)

        module.call = call
        precompute = Precomputation(module=module, input_domain=torch.tensor([[[1]]]))
        precompute()
        actual = json.dumps(precompute, cls=JSONEncoder)
        expected = """{"description": [], "shape": [1, 1], "x": [[[1]]], "y": [3]}"""
        self.assertEqual(expected, actual)

    def test_precomputable_tag(self):
        input_generator = [-1, 1]

        @precomputable(input_shape=[-1, 1], input_generator=input_generator)
        class MyCustomModule(torch.nn.Module):
            def __init__(self):
                super().__init__()

        module = MyCustomModule()
        tags = get_tags(module)

        self.assertSequenceEqual(tags["precomputable"]["input_shape"], [-1, 1])
