# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE.pytorch for license information.

import typing

import torch
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | FileCheck %s

mb = ModuleBuilder()


class Submodule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.exported = 1
        self.not_exported = 2

    def forward(self):
        return self.not_exported_method()

    def not_exported_method(self):
        return


class TestModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.s = Submodule()

    def forward(self, tensor, value_tensor):
        return self.s.forward()


test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)

annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()

annotator.exportNone(class_type)
annotator.exportPath(class_type, ['s', 'exported'])
annotator.exportPath(class_type, ['s', 'forward'])
annotator.annotateArgs(class_type, ['forward'], [
    None,
    ((1024, 2), torch.float32, False),
    ((42, -1, 7), torch.int8, True),
])

# "Change detector" test + "documentation" for the repr of `ClassAnnotator`.
# This is semi-load-bearing because users interact with this class and repr
# will show up in error messages, so should be pretty readable.
#
# CHECK:      ClassAnnotator {
# CHECK-NEXT:   ClassAnnotation('__torch__.Submodule') {
# CHECK-NEXT:     AttributeAnnotation('training') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:     }
# CHECK-NEXT:     AttributeAnnotation('_is_full_backward_hook') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:     }
# CHECK-NEXT:     AttributeAnnotation('exported') {
# CHECK-NEXT:       isExported = true
# CHECK-NEXT:     }
# CHECK-NEXT:     AttributeAnnotation('not_exported') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:     }
# CHECK-NEXT:     MethodAnnotation('forward') {
# CHECK-NEXT:       isExported = true
# CHECK-NEXT:       argAnnotations = <none>
# CHECK-NEXT:     }
# CHECK-NEXT:     MethodAnnotation('not_exported_method') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:       argAnnotations = <none>
# CHECK-NEXT:     }
# CHECK-NEXT:   }
# CHECK-NEXT:   ClassAnnotation('__torch__.TestModule') {
# CHECK-NEXT:     AttributeAnnotation('training') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:     }
# CHECK-NEXT:     AttributeAnnotation('_is_full_backward_hook') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:     }
# CHECK-NEXT:     AttributeAnnotation('s') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:     }
# CHECK-NEXT:     MethodAnnotation('forward') {
# CHECK-NEXT:       isExported = false
# CHECK-NEXT:       argAnnotations =
# CHECK-NEXT:         ArgAnnotation(0) {
# CHECK-NEXT:           dtype = <none>
# CHECK-NEXT:           shape = <none>
# CHECK-NEXT:           hasValueSemantics = false
# CHECK-NEXT:         }
# CHECK-NEXT:         ArgAnnotation(1) {
# CHECK-NEXT:           dtype = Float
# CHECK-NEXT:           shape = [1024, 2]
# CHECK-NEXT:           hasValueSemantics = false
# CHECK-NEXT:         }
# CHECK-NEXT:         ArgAnnotation(2) {
# CHECK-NEXT:           dtype = Char
# CHECK-NEXT:           shape = [42, -1, 7]
# CHECK-NEXT:           hasValueSemantics = true
# CHECK-NEXT:         }
# CHECK-NEXT:     }
# CHECK-NEXT:   }
# CHECK-NEXT: }

print(annotator)
