#   Copyright (c) 2020 PaddlePaddle 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 paddle.base import unique_name
from paddle.jit.dy2static.utils import (
    ORIGIN_INFO,
    ast_to_source_code,
)
from paddle.utils import gast

from .utils import (
    FOR_ITER_INDEX_PREFIX,
    FOR_ITER_ITERATOR_PREFIX,
    FOR_ITER_TARGET_PREFIX,
    FOR_ITER_VAR_LEN_PREFIX,
    FOR_ITER_VAR_NAME_PREFIX,
    FOR_ITER_ZIP_TO_LIST_PREFIX,
    create_assign_node,
)

__all__ = []


class BaseTransformer(gast.NodeTransformer):
    def visit(self, node):
        if not isinstance(node, gast.AST):
            msg = f'Expected "gast.AST", but got "{type(node)}".'
            raise ValueError(msg)
        origin_info = getattr(node, ORIGIN_INFO, None)

        result = super().visit(node)

        iter_result = result
        if iter_result is not node and iter_result is not None:
            if not isinstance(iter_result, (list, tuple)):
                iter_result = (iter_result,)
            if origin_info is not None:
                for n in iter_result:
                    setattr(n, ORIGIN_INFO, origin_info)

        return result


class NameNodeReplaceTransformer(BaseTransformer):
    """
    This class replaces specified gast.Name node by replace_node.
    """

    def __init__(self, root_node, target_name, replace_node):
        assert isinstance(target_name, str)

        # NOTE(liym27):
        # Use gast.Name to replace gast.Name, otherwise, errors may occur.
        #
        # For examples:
        # If using a gast.Subscript to replace gast.Name, and the original gast.Name
        # is in the arguments of FunctionDef, an exception will be raised.
        #
        # ```
        # def func(x[i])) # x[i] can not be a argument
        #    # ...
        # ```

        assert isinstance(replace_node, gast.Name)
        self.target_name = target_name
        self.replace_node = replace_node

        self.visit(root_node)

    def visit_Name(self, node):
        if node.id == self.target_name:
            return self.replace_node
        return node

    def visit_Nonlocal(self, node):
        names = node.names

        def replace(s):
            if s == self.target_name:
                return self.replace_node.id
            return s

        node.names = list(map(replace, names))
        return node


class ForLoopTuplePreTransformer(BaseTransformer):
    """pre-process of for loop.
    >>> for A in B:
    >>>    C

    will be changed into :

    >>> # make iterator-only to indexable list.
    >>> UUID_iterator = _jst.Indexable(B)
    >>> for UUID_target in UUID_iterator:
    >>>     A = _jst.Unpack(UUID_target, structure)
    >>>     C

    make the later loop_transform have unified type:
    >>> for target in iter:
    >>>     body
    """

    def __init__(self, root):
        self.root = root

    def transform(self):
        self.visit(self.root)

    def visit_For(self, node):
        self.generic_visit(node)
        tuple_target = unique_name.generate(FOR_ITER_TARGET_PREFIX)
        tuple_iterator = unique_name.generate(FOR_ITER_ITERATOR_PREFIX)
        origin_tuple_node = node.target
        assign_iterator_node = gast.parse(
            f"{tuple_iterator} = _jst.Indexable({ast_to_source_code(node.iter).strip()})"
        ).body[0]
        node.target = gast.Name(
            id=tuple_target,
            ctx=gast.Store(),
            annotation=None,
            type_comment=None,
        )
        node.iter = gast.Name(
            id=tuple_iterator,
            ctx=gast.Load(),
            annotation=None,
            type_comment=None,
        )
        node.body[0:0] = self.tuple_to_stmts(origin_tuple_node, tuple_target)
        # return a list will insert a list of node replace the original for node.
        return [assign_iterator_node, node]

    def tuple_node_to_unpack_structure(self, node):
        """Create a sequence to represents the structure of nest.
        For example: `a, (b,c), [d,e,f]` is represented by
        `[1, [1,1], [1,1,1]]`. the `1` is just a notation.

        Specially, `a` is represented by `1`.
        """
        ret = []
        if not isinstance(node, (gast.Tuple, gast.List)):
            return 1
        for element in node.elts:
            ret.append(self.tuple_node_to_unpack_structure(element))
        return ret

    def tuple_to_stmts(self, node, tuple_name):
        structure_str = str(self.tuple_node_to_unpack_structure(node))
        node_str = ast_to_source_code(node).strip()
        assign_node_str = (
            f"{node_str} = _jst.Unpack({tuple_name}, {structure_str})"
        )
        assign_node = gast.parse(assign_node_str).body[0]
        return [assign_node]


class ForNodeVisitor:
    """
    This class parses python for statement, get transformed 3 statement components of for node
    three key statements:
        1). init_stmts: list[node], prepare nodes of for loop, may not only one
        2). cond_stmt: node, condition node to judge whether continue loop
        3). body_stmts: list[node], updated loop body, sometimes we should change
            the original statement in body, not just append new statement

    In this process, the semantics of for does not change.

    Now only can parse 3 type statements (Here var is Tensor(Tensor) or python variable):
        1). for x in range(var[*]|var.numpy()[*])
        2). for x in var|var.numpy()
        3). for i, x enumerate(var|var.numpy())
    """

    def __init__(self, for_node):
        assert isinstance(for_node, gast.For), (
            "Input node for the initialization of ForNodeVisitor is not gast.For node."
        )
        # 1. original for node
        self.node = for_node

        # 2. gast.For node main parts
        self.target = for_node.target
        # NOTE: type may be Node or list[Node]
        self.iter_args = (
            for_node.iter if self.is_for_iter() else for_node.iter.args
        )
        self.body = for_node.body

        # 3. key shared node or names
        # - x:
        #   - for x in range(***)
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
        self.iter_var_name = self._get_iter_var_name()

        # - created index var to slice Variable: __for_loop_var_index_0
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
        self.iter_idx_name = unique_name.generate(FOR_ITER_INDEX_PREFIX)

        # - created shape var to build loop condition: __for_loop_var_len_0
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
        #   - for x in var
        self.iter_var_len_name = unique_name.generate(FOR_ITER_VAR_LEN_PREFIX)
        # - created zip to list var : __for_loop_iter_zip_0
        self.iter_zip_to_list_name = unique_name.generate(
            FOR_ITER_ZIP_TO_LIST_PREFIX
        )

        # - var.numpy()/var
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
        self.iter_node = self._get_iter_node()

        # - enumerate i:
        #   - for i, x enumerate(var|var.numpy())
        self.enum_idx_name = self._get_enum_idx_name()

        # - range/enumerate args length
        self.args_length = None

    def parse(self):
        self._args_check()
        if self.is_for_range_iter():
            return self._parse_for_range_stmts()
        elif self.is_for_iter():
            return self._parse_for_stmts()
        elif self.is_for_enumerate_iter():
            return self._parse_for_enumerate_stmts()
        else:
            return None

    def is_for_range_iter(self):
        return (
            isinstance(self.node.iter, gast.Call)
            and isinstance(self.node.iter.func, gast.Name)
            and self.node.iter.func.id == "range"
        )

    def is_for_iter(self):
        if isinstance(
            self.node.iter, (gast.Name, gast.Attribute, gast.List, gast.Tuple)
        ):
            return True
        elif (
            isinstance(self.node.iter, gast.Call)
            and isinstance(self.node.iter.func, gast.Attribute)
            and self.node.iter.func.attr == 'numpy'
        ):
            return True
        elif isinstance(self.node.iter, gast.Subscript):
            return True
        else:
            return False

    def is_for_enumerate_iter(self):
        return (
            isinstance(self.node.iter, gast.Call)
            and isinstance(self.node.iter.func, gast.Name)
            and self.node.iter.func.id == "enumerate"
        )

    def _args_check(self):
        if self.is_for_range_iter():
            self.args_length = len(self.iter_args)
            assert self.args_length >= 1 and self.args_length <= 3, (
                "range() function takes 1 to 3 arguments"
            )
        elif self.is_for_enumerate_iter():
            self.args_length = len(self.iter_args)
            assert self.args_length >= 1 and self.args_length <= 2, (
                "enumerate() function takes 1 to 2 arguments"
            )
        else:
            self.args_length = None

    def _parse_for_range_stmts(self):
        init_stmts = []
        init_stmts.append(self._build_index_init_node())

        compare_node = self._build_compare_node()
        step_node = self._build_step_node()
        cond_stmt = self._build_cond_stmt(step_node, compare_node)

        body_stmts = self.body
        body_stmts.append(self._build_index_increase_node(step_node))

        return init_stmts, cond_stmt, body_stmts

    def _parse_for_stmts(self):
        init_stmts = []
        init_stmts.extend(self._build_iter_node())
        init_stmts.append(self._build_index_init_node())
        init_stmts.append(self._build_var_len_assign_node())

        compare_node = self._build_compare_node()
        step_node = self._build_step_node()
        cond_stmt = self._build_cond_stmt(step_node, compare_node)

        body_stmts = self.body

        # NOTE(liym27): Here add a gast.Assign, and the target of it is gast.Name.
        # In NameNodeReplaceTransformer, using gast.Name to replace gast.Name is safe.
        target_node, assign_node = self._build_assign_var_slice_node()
        body_stmts[0:0] = [assign_node]
        for body_node in body_stmts:
            NameNodeReplaceTransformer(
                body_node, self.iter_var_name, target_node
            )
        body_stmts.append(self._build_index_increase_node(step_node))

        return init_stmts, cond_stmt, body_stmts

    def _parse_for_enumerate_stmts(self):
        init_stmts = []
        init_stmts.extend(self._build_iter_node())
        init_stmts.append(self._build_index_init_node())
        init_stmts.append(self._build_var_len_assign_node())
        init_stmts.append(self._build_enum_init_node())

        compare_node = self._build_compare_node()
        step_node = self._build_step_node()
        cond_stmt = self._build_cond_stmt(step_node, compare_node)

        body_stmts = self.body

        target_node, assign_node = self._build_assign_var_slice_node()
        body_stmts[0:0] = [assign_node]
        for body_node in body_stmts:
            NameNodeReplaceTransformer(
                body_node, self.iter_var_name, target_node
            )

        body_stmts.append(self._build_index_increase_node(step_node))
        body_stmts.append(self._build_enum_increase_node())

        return init_stmts, cond_stmt, body_stmts

    def _build_index_init_node(self):
        if self.is_for_range_iter():
            if self.args_length == 1:
                index_init_value_str = '0'
            else:
                index_init_value_str = ast_to_source_code(
                    self.iter_args[0]
                ).strip()

            index_init_var_name = self.iter_var_name
        else:
            index_init_value_str = '0'
            index_init_var_name = self.iter_idx_name

        index_init_node_source_str = (
            f"{index_init_var_name} = {index_init_value_str}"
        )

        index_init_node = gast.parse(index_init_node_source_str).body[0]

        return index_init_node

    def _build_var_len_assign_node(self):
        # get the length of iterable variable
        if (
            isinstance(self.iter_node, gast.Call)
            and isinstance(self.iter_node.func, gast.Attribute)
            and self.iter_node.func.attr == 'numpy'
        ):
            iter_var_name = ast_to_source_code(
                self.iter_node.func.value
            ).strip()
        else:
            iter_var_name = ast_to_source_code(self.iter_node).strip()

        convert_len_node_source_str = (
            f'{self.iter_var_len_name} = _jst.Len({iter_var_name})'
        )

        convert_len_node = gast.parse(convert_len_node_source_str).body[0]

        return convert_len_node

    def _build_iter_node(self):
        """
        Process special cases for iter_node include:
          - Case 1 (for zip):

            - for i, val in enumerate(zip(x, y))  # original code:

            - __for_loop_iter_zip_0 = list(zip(x, y))
            - for i, val in enumerate(__for_loop_iter_zip_0)
        """
        new_nodes = []
        if isinstance(self.iter_node, gast.Call) and isinstance(
            self.iter_node.func, gast.Name
        ):
            if self.iter_node.func.id == 'zip':
                iter_var_name = ast_to_source_code(self.iter_node).strip()
                zip_to_list_str = (
                    f"{self.iter_zip_to_list_name} = list({iter_var_name})"
                )
                zip_to_list_node = gast.parse(zip_to_list_str).body[0]
                new_nodes.append(zip_to_list_node)

                self.iter_node = gast.Name(
                    id=self.iter_zip_to_list_name,
                    ctx=gast.Load(),
                    annotation=None,
                    type_comment=None,
                )

        return new_nodes

    def _build_enum_init_node(self):
        if self.is_for_enumerate_iter() and self.args_length != 1:
            init_value_str = ast_to_source_code(self.iter_args[1]).strip()
        else:
            init_value_str = '0'

        enum_init_node_source_str = f"{self.enum_idx_name} = {init_value_str}"
        enum_init_node = gast.parse(enum_init_node_source_str).body[0]
        return enum_init_node

    def _build_compare_node(self):
        if self.is_for_range_iter():
            compare_node = (
                self.iter_args[0]
                if self.args_length == 1
                else self.iter_args[1]
            )
        else:
            compare_node = gast.Name(
                id=self.iter_var_len_name,
                ctx=gast.Load(),
                annotation=None,
                type_comment=None,
            )
        return compare_node

    def _build_step_node(self):
        if self.is_for_range_iter():
            step_node = (
                self.iter_args[2]
                if self.args_length == 3
                else gast.Constant(value=1, kind=None)
            )
        else:
            step_node = gast.Constant(value=1, kind=None)
        return step_node

    def _build_cond_stmt(self, step_node, compare_node):
        if not isinstance(step_node, (gast.Constant, gast.UnaryOp)):
            raise NotImplementedError(
                "Dynamic-to-Static only supports the step value is a constant or negative constant in 'for-range' statements, "
                f"such as '2', '-3'. But received: '{ast_to_source_code(step_node).strip()}'. Please fix code to be compatible with Dynamic-to-Static."
            )

        if isinstance(step_node, gast.UnaryOp) or step_node.value < 0:
            # eg:
            # range(max, min, -2)
            # ->
            # i > min
            return gast.Compare(
                left=gast.Name(
                    id=(
                        self.iter_var_name
                        if self.is_for_range_iter()
                        else self.iter_idx_name
                    ),
                    ctx=gast.Load(),
                    annotation=None,
                    type_comment=None,
                ),
                ops=[gast.Gt()],
                comparators=[compare_node],
            )
        else:
            # eg:
            # range(min, max, 2)
            # ->
            # i < max
            return gast.Compare(
                left=gast.Name(
                    id=(
                        self.iter_var_name
                        if self.is_for_range_iter()
                        else self.iter_idx_name
                    ),
                    ctx=gast.Load(),
                    annotation=None,
                    type_comment=None,
                ),
                ops=[gast.Lt()],
                comparators=[compare_node],
            )

    def _build_index_increase_node(self, step_node):
        return gast.AugAssign(
            target=gast.Name(
                id=(
                    self.iter_var_name
                    if self.is_for_range_iter()
                    else self.iter_idx_name
                ),
                ctx=gast.Store(),
                annotation=None,
                type_comment=None,
            ),
            op=gast.Add(),
            value=step_node,
        )

    def _build_assign_var_slice_node(self):
        var_slice_str = f"{ast_to_source_code(self.iter_node).strip()}[{self.iter_idx_name}]"
        var_slice_node = gast.parse(var_slice_str).body[0].value
        new_iter_var_name = unique_name.generate(FOR_ITER_VAR_NAME_PREFIX)
        target_node, assign_node = create_assign_node(
            new_iter_var_name, var_slice_node
        )
        return target_node, assign_node

    def _build_enum_increase_node(self):
        return gast.AugAssign(
            target=gast.Name(
                id=self.enum_idx_name,
                ctx=gast.Store(),
                annotation=None,
                type_comment=None,
            ),
            op=gast.Add(),
            value=gast.Constant(value=1, kind=None),
        )

    def _get_iter_var_name(self):
        if self.is_for_range_iter():
            return self.target.id
        elif self.is_for_iter():
            return self.target.id
        elif self.is_for_enumerate_iter():
            return self.target.elts[1].id
        return None

    def _get_iter_node(self):
        if self.is_for_iter():
            return self.iter_args
        elif self.is_for_enumerate_iter():
            return self.iter_args[0]
        return None

    def _get_enum_idx_name(self):
        if self.is_for_enumerate_iter():
            return self.target.elts[0].id
        return None