"""
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from collections import OrderedDict
from functools import partial
from typing import Iterator, Optional, Set, Union, List

import oneflow._C
import oneflow._oneflow_internal
import oneflow.framework.graph_build_util as graph_build_util
from oneflow.env import get_rank
from oneflow.framework.tensor import Tensor, TensorTuple
from oneflow.nn.module import Module
from oneflow.nn.modules.container import *
from oneflow.nn.utils.container import *
from oneflow.nn.parameter import Parameter
from oneflow.nn.graph.block_config import BlockConfig
from oneflow.nn.graph.util import add_indent, seq_to_func_return


def get_block_cls(item):
    if isinstance(item, Sequential):
        return SequentialBlock
    elif isinstance(item, ModuleList):
        return ModuleListBlock
    elif isinstance(item, ModuleDict):
        return ModuleDictBlock
    elif isinstance(item, ParameterList):
        return ParameterListBlock
    elif isinstance(item, ParameterDict):
        return ParameterDictBlock
    elif isinstance(item, Module):
        return ModuleBlock
    elif isinstance(item, Tensor):
        return TensorBlock
    else:
        raise NotImplementedError()


class BlockType:
    NONE = "NONE"
    MODULE = "MODULE"
    PARAMETER = "PARAMETER"
    BUFFER = "BUFFER"


class Block(object):
    def __init__(
        self, prefix: str = "", name: str = "",
    ):
        self._name = name
        self._name_prefix = prefix
        self._type = BlockType.NONE
        self._origin = None
        self._scope = None
        self._prev_scope = None
        self.config = BlockConfig()

    @property
    def name(self):
        return self._name

    @property
    def name_prefix(self):
        return self._name_prefix

    @property
    def type(self):
        return self._type

    @property
    def prev_scope(self):
        if self._prev_scope is None:
            self._prev_scope = oneflow._oneflow_internal.GetCurrentScope()
        return self._prev_scope

    @property
    def scope(self):
        if self._scope is None:
            self._scope = graph_build_util.make_new_block_scope(self.prev_scope, self)
        return self._scope

    def scope_context(self):
        return graph_build_util.BlockScopeContext(self.prev_scope, self.scope)


class ModuleBlock(Block):
    def __init__(
        self, prefix: str = "", name: str = "", origin: Module = None,
    ):
        assert not isinstance(origin, Block)
        super().__init__(prefix, name)
        self._debug = False
        self._debug_min_s_level = 2
        self._debug_max_v_level = 0
        self._type = BlockType.MODULE
        self._is_executing_forward = False
        self._modules = OrderedDict()
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._args_repr = []
        self._outs_repr = []
        self.set_origin(origin)

    @property
    def origin(self):
        return self._origin

    def set_origin(self, origin):
        self._origin = origin
        if origin is None:
            return
        assert isinstance(origin, Module)
        for (n, m) in list(origin.named_children()):
            self.__setattr__(
                n, get_block_cls(m)(self._name_prefix + self._name + ".", n, m)
            )
        for (n, p) in list(origin.named_parameters("", False)):
            self.__setattr__(
                n, get_block_cls(p)(self._name_prefix + self._name + ".", n, p)
            )
        for (n, b) in list(origin.named_buffers("", False)):
            self.__setattr__(
                n, get_block_cls(b)(self._name_prefix + self._name + ".", n, b)
            )

    def debug(
        self,
        v_level: int = 0,
        ranks: Optional[Union[int, List[int]]] = None,
        mode: bool = True,
    ) -> None:
        assert isinstance(mode, bool)
        assert isinstance(v_level, int)

        if ranks is None:
            rank_list = [0]
        elif isinstance(ranks, int):
            rank_list = [ranks]
        elif isinstance(ranks, list):
            rank_list = ranks
        else:
            raise ValueError("ranks must be int or List[int].")

        my_rank = get_rank()
        if -1 in rank_list or my_rank in rank_list:
            self._debug = mode
            if self._debug:
                self._debug_min_s_level = 0
                self._debug_max_v_level = v_level
            if self._type == BlockType.MODULE:

                def _set_child(d):
                    for (_, n) in d.items():
                        n.debug(v_level, ranks, mode)

                _set_child(self._modules)

    def __call__(self, *args):
        assert self._type == BlockType.MODULE
        self._print(0, 1, self._shallow_repr())

        for idx, arg in enumerate(args):
            meta_repr_str = (
                arg._meta_repr() if isinstance(arg, Tensor) else str(type(arg))
            )
            in_str = (
                "(INPUT:_"
                + self.name_prefix
                + self.name
                + "-input_"
                + str(idx)
                + ":"
                + meta_repr_str
                + ")"
            )
            if not isinstance(arg, Tensor):
                in_str = "[WARNING]" + in_str
            self._args_repr.append(in_str)

            self._print(0, 1, in_str)

            def _print_state(d):
                for (_, n) in d.items():
                    self._print(0, 1, n._shallow_repr())

            _print_state(self._parameters)
            _print_state(self._buffers)

        # NOTE: The original nn.Moudle's __call__ method is ignored, which means
        # that hooks of nn.Modules are ignored. It is not recommended
        # to use hooks of nn.Module in nn.Graph for the moment.
        # result = self._origin.__class__.__call__(self, *args)
        result = self._forward(*args)

        outputs = ()
        if not (type(result) is tuple or type(result) is list):
            outputs = (result,)
        else:
            outputs = result

        for idx, out in enumerate(outputs):
            out_repr = out._meta_repr() if isinstance(out, Tensor) else str(type(out))
            out_str = (
                "(OUTPUT:_"
                + self.name_prefix
                + self.name
                + "-output_"
                + str(idx)
                + ":"
                + out_repr
                + ")"
            )
            if not isinstance(out, Tensor):
                out_str = "[WARNING]" + out_str

            self._outs_repr.append(out_str)
            self._print(0, 1, out_str)

        return result

    def _forward(self, *args):
        self._is_executing_forward = True
        args = self._pre_forward_mapping_out_scope(*args)
        with self.scope_context():
            result = self._origin.__class__.forward(self, *args)
        result = self._post_forward_mapping_out_scope(result)
        result = seq_to_func_return(result)
        self._is_executing_forward = False
        return result

    def _pre_forward_mapping_out_scope(self, *args):
        # Insert identity op when doing activation checkpointing or pipeline execution.
        # Identity op outside activation checkpointing scope will be the endpoint of an activation checkpointing segment.
        # Identity op as the first op of a pipeline stage will make backward op depends on the identity op within the stage,
        # otherwise the backward op may depends the op in former stage which will make graph creates unnessary buffers.
        if self.config.activation_checkpointing or (
            self.config.stage_id is not None and self.config.stage_id >= 0
        ):

            def insert_identity(t):
                assert isinstance(t, Tensor)
                return oneflow._C.identity(t)

            args = self._mapping_io("input", insert_identity, "insert_identity", *args,)

        return args

    def _post_forward_mapping_out_scope(self, *args):
        # Insert identity op when doing activation checkpointing or pipeline execution.
        if self.config.activation_checkpointing or (
            self.config.stage_id is not None and self.config.stage_id >= 0
        ):

            def insert_identity(t):
                assert isinstance(t, Tensor)
                return oneflow._C.identity(t)

            args = self._mapping_io(
                "output", insert_identity, "insert_identity", *args,
            )
        return args

    def add_module(self, name: str, module: Optional[Module]) -> None:
        self.__setattr__(
            name,
            get_block_cls(module)(self._name_prefix + self._name + ".", name, module),
        )

    def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
        self.__setattr__(
            name,
            get_block_cls(param)(self._name_prefix + self._name + ".", name, param),
        )

    def modules(self, memo: Optional[Set["Block"]] = None) -> Iterator["Block"]:
        assert self._type == BlockType.MODULE
        if memo is None:
            memo = set()
        if self not in memo:
            memo.add(self)
            yield self
            for (name, module) in self._modules.items():
                if module is None:
                    continue
                for m in module.modules(memo):
                    yield m

    def _mapping_io(self, io_type, func, func_desc, *args):
        assert isinstance(func_desc, str)
        assert io_type in ("input", "output")
        mapped_args = []

        def mapping_tensor(item):
            assert isinstance(item, Tensor)
            return func(item)

        for idx, arg in enumerate(args):
            if isinstance(arg, list):
                seq_args = list()
                for i in range(len(arg)):
                    is_tensor, name, repr_str = self._io_tensor_check_and_gen(
                        arg[i], io_type, idx, i
                    )
                    if is_tensor:
                        seq_args.append(mapping_tensor(arg[i]))
                        self._print(
                            0,
                            1,
                            f"{repr_str} is a Tensor, {func_desc} transformation has been done.",
                        )
                    else:
                        self._print(
                            0,
                            0,
                            f"{repr_str} is not a Tensor, {func_desc} transformation will be ignored.",
                        )
                        seq_args.append(arg[i])
                mapped_args.append(seq_args)
            elif isinstance(arg, Tensor):
                is_tensor, name, repr_str = self._io_tensor_check_and_gen(
                    arg, io_type, idx
                )
                assert is_tensor
                mapped_args.append(mapping_tensor(arg))
                self._print(
                    0,
                    1,
                    f"{repr_str} is a Tensor, {func_desc} transformation has been done.",
                )
            else:
                is_tensor, name, repr_str = self._io_tensor_check_and_gen(
                    arg, io_type, idx
                )
                assert not is_tensor
                mapped_args.append(arg)
                self._print(
                    0,
                    0,
                    f"{repr_str} is not a Tensor or a list of Tensor, {func_desc} transformation will be ignored.",
                )

        return tuple(mapped_args)

    def _io_tensor_check_and_gen(self, item, io_type, idx, second_idx=None):
        assert io_type in ("input", "output")
        name = (
            "_"
            + self.name_prefix
            + self.name
            + "-"
            + io_type
            + "_"
            + str(idx)
            + ("" if second_idx is None else "_" + str(second_idx))
        )
        if isinstance(item, Tensor):
            repr_str = (
                "(" + io_type.upper() + ":" + name + ":" + item._meta_repr() + ")"
            )
            return True, name, repr_str
        else:
            repr_str = (
                "[WARNING]("
                + io_type.upper()
                + ":"
                + name
                + ":"
                + str(type(item))
                + ")"
            )
            return False, name, repr_str

    def _members(self, get_members_fn, recurse=True) -> Iterator["Block"]:
        assert self._type == BlockType.MODULE
        memo = set()
        modules = self.modules() if recurse else [self]
        for module in modules:
            members = get_members_fn(module)
            for (k, v) in members:
                if v is None or v in memo:
                    continue
                memo.add(v)
                yield v

    def parameters(self, recurse: bool = True) -> Iterator["Block"]:
        assert self._type == BlockType.MODULE
        gen = self._members(lambda module: module._parameters.items(), recurse=recurse)
        for elem in gen:
            yield elem

    def buffers(self, recurse: bool = True) -> Iterator["Block"]:
        assert self._type == BlockType.MODULE
        gen = self._members(lambda module: module._buffers.items(), recurse=recurse)
        for elem in gen:
            yield elem

    def __setattr__(self, name: str, value=None) -> None:
        if value is None or not isinstance(value, Block):
            self.__dict__[name] = value
        else:
            dicts_or_sets = (
                self.__dict__,
                self._modules,
                self._parameters,
                self._buffers,
            )
            for d in dicts_or_sets:
                if name in d:
                    raise AttributeError(
                        "'{}' object has duplicated attribute named '{}'".format(
                            self._name, name
                        )
                    )
            if value.type == BlockType.MODULE:
                self._modules[name] = value
            elif value.type == BlockType.PARAMETER:
                self._parameters[name] = value
            elif value.type == BlockType.BUFFER:
                self._buffers[name] = value
            else:
                raise AttributeError(
                    "'{}' object are not allowed to set attribute named '{}'".format(
                        type(self).__name__, name
                    )
                )

    def __getattr__(self, name: str):
        if name in self.__dict__:
            return self.__dict__[name]
        # support get module
        if "_modules" in self.__dict__:
            modules = self.__dict__["_modules"]
            if name in modules:
                return modules[name]
        # support get parameter
        p_state = self._get_from_states(name, "_parameters")
        if p_state is not None:
            return p_state
        # support get buffer
        b_state = self._get_from_states(name, "_buffers")
        if b_state is not None:
            return b_state
        # support get normal attr
        if name in self._origin.__dict__:
            return self._origin.__dict__[name]
        # support get function
        if hasattr(self._origin, name):
            return partial(getattr(self._origin.__class__, name), self)
        raise AttributeError(
            "'{}' '{}' object '{}' in nn.Graph has no attribute '{}'".format(
                self._type, type(self).__name__, self._name_prefix + self.name, name
            )
        )

    def _get_from_states(self, name, states_name):
        if states_name not in self.__dict__:
            return None

        _states = self.__dict__[states_name]
        if name not in _states:
            return None

        _s_block = _states[name]
        if graph_build_util.lazy_mode.is_enabled():
            #  lazy
            if _s_block._lazy_origin is None:
                assert _s_block._lazy_origin_builder is not None, (
                    repr(_s_block) + " has no lazy Tensor creation function."
                )
                assert self._is_executing_forward, (
                    repr(_s_block)
                    + "'s first get must happened in it's nn.Module.forward() to generate the right scope."
                )
                with _s_block.scope_context():
                    _s_block._lazy_origin = _s_block._lazy_origin_builder()
            return _s_block._lazy_origin
        elif (
            not graph_build_util.lazy_mode.is_enabled()
        ) and self._is_executing_forward:
            # eager and inside nn.Graph.build()
            return _s_block.origin
        else:
            # outside nn.Graph.build()
            return _s_block

    def __repr__(self):
        lines = None
        child_lines = []
        if (self.config is not None) and (not self.config._is_null):
            child_lines.append(add_indent(repr(self.config), 2))
        if len(self._args_repr) > 0:
            for in_str in self._args_repr:
                input_str = add_indent(in_str, 2)
                child_lines.append(input_str)

        def _append_child(d):
            for (_, n) in d.items():
                n_str = repr(n)
                n_str = add_indent(n_str, 2)
                child_lines.append(n_str)

        _append_child(self._parameters)
        _append_child(self._buffers)
        _append_child(self._modules)

        if len(self._outs_repr) > 0:
            for out_str in self._outs_repr:
                output_str = add_indent(out_str, 2)
                child_lines.append(output_str)

        if len(child_lines) > 0:
            lines = child_lines

        main_str = self._shallow_repr() + ": ("
        if lines is not None:
            main_str += "\n  " + "\n  ".join(lines) + "\n"
        main_str += ")"
        return main_str

    def _shallow_repr(self):
        shallow_repr = (
            "("
            + self._type
            + ":"
            + self._name_prefix
            + self._name
            + ":"
            + self._origin._shallow_repr()
            + ")"
        )
        return shallow_repr

    def _print(self, s_level=2, v_level=0, msg: str = ""):
        r"""Do print according to info level.
        """
        assert isinstance(s_level, int)
        assert isinstance(v_level, int)
        assert isinstance(msg, str)
        if s_level >= self._debug_min_s_level:
            if (s_level > 0) or (s_level == 0 and v_level <= self._debug_max_v_level):
                print(msg)


class TensorBlock(Block):
    def __init__(
        self, prefix: str = "", name: str = "", origin: Union[Parameter, Tensor] = None,
    ):
        assert not isinstance(origin, Block)
        super().__init__(prefix, name)
        if isinstance(origin, Parameter):
            self._type = BlockType.PARAMETER
        elif isinstance(origin, Tensor):
            self._type = BlockType.BUFFER
        else:
            raise NotImplementedError()
        self._lazy_origin = None
        self._lazy_origin_builder = None
        self.set_origin(origin)

    @property
    def origin(self):
        return self._origin

    def set_origin(self, origin):
        self._origin = origin

    @property
    def lazy_origin(self):
        assert (
            self._type == BlockType.PARAMETER or self._type == BlockType.BUFFER
        ), "Only Parameter or Buffer Block has lazy_origin"
        return self._lazy_origin

    def lazy_origin_builder(self):
        assert (
            self._type == BlockType.PARAMETER or self._type == BlockType.BUFFER
        ), "Only Parameter or Buffer Block has lazy_origin_builder"
        return self._lazy_origin_builder

    def set_lazy_origin_builder(self, fn=None):
        assert (
            self._type == BlockType.PARAMETER or self._type == BlockType.BUFFER
        ), "Only Parameter or Buffer Block has lazy_origin_builder"
        self._lazy_origin_builder = fn

    def __repr__(self):
        lines = None
        main_str = self._shallow_repr() + ": ("
        if lines is not None:
            main_str += "\n  " + "\n  ".join(lines) + "\n"
        main_str += ")"
        return main_str

    def _shallow_repr(self):
        shallow_repr = (
            "("
            + self._type
            + ":"
            + self._name_prefix
            + self._name
            + ":"
            + self._origin._meta_repr()
            + ")"
        )
        return shallow_repr


class SequentialBlock(get_seq(ModuleBlock)):
    def __init__(
        self, prefix: str = "", name: str = "", origin: Sequential = None,
    ):
        super().__init__()
        self._name_prefix = prefix
        self._name = name
        self.set_origin(origin)


class ModuleListBlock(get_list(ModuleBlock)):
    def __init__(
        self, prefix: str = "", name: str = "", origin: ModuleList = None,
    ):
        super().__init__()
        self._name_prefix = prefix
        self._name = name
        self.set_origin(origin)
        # MoudleList is a container without forward() method,
        # so it will not be executed or has an execution config.
        self.config = None


class ModuleDictBlock(get_dict(ModuleBlock)):
    def __init__(
        self, prefix: str = "", name: str = "", origin: ModuleDict = None,
    ):
        super().__init__()
        self._name_prefix = prefix
        self._name = name
        self.set_origin(origin)


class ParameterListBlock(get_para_list(ModuleBlock)):
    def __init__(
        self, prefix: str = "", name: str = "", origin: ParameterList = None,
    ):
        super().__init__()
        self._name_prefix = prefix
        self._name = name
        self.set_origin(origin)
        self._is_executing_forward = True

    def __getitem__(self, idx):
        assert isinstance(idx, int)
        idx = self._get_abs_string_index(idx)
        key = str(idx)
        p_state = self._get_from_states(key, "_parameters")
        if p_state is not None:
            return p_state
        else:
            raise AttributeError("ParameterList dosen't contain ", key)


class ParameterDictBlock(get_para_dict(ModuleBlock)):
    def __init__(
        self, prefix: str = "", name: str = "", origin: ParameterDict = None,
    ):
        super().__init__()
        self._name_prefix = prefix
        self._name = name
        self.set_origin(origin)
        self._is_executing_forward = True

    def __getitem__(self, key: str):
        p_state = self._get_from_states(key, "_parameters")
        if p_state is not None:
            return p_state
        else:
            raise AttributeError("ParameterDict dosen't contain key ", key)
