#   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 __future__ import print_function

import six
import copy
from collections import defaultdict

# gast is a generic AST to represent Python2 and Python3's Abstract Syntax Tree(AST).
# It provides a compatibility layer between the AST of various Python versions,
# as produced by ast.parse from the standard ast module.
# See details in https://github.com/serge-sans-paille/gast/
import gast
from paddle.fluid import unique_name

from paddle.fluid.dygraph.dygraph_to_static.utils import create_funcDef_node, ast_to_source_code
from paddle.fluid.dygraph.dygraph_to_static.utils import create_assign_node
from paddle.fluid.dygraph.dygraph_to_static.static_analysis import StaticAnalysisVisitor
from paddle.fluid.dygraph.dygraph_to_static.static_analysis import AstNodeWrapper
from paddle.fluid.dygraph.dygraph_to_static.variable_trans_func import create_static_variable_gast_node

TRUE_FUNC_PREFIX = 'true_fn'
FALSE_FUNC_PREFIX = 'false_fn'


class IfElseTransformer(gast.NodeTransformer):
    """
    Transform if/else statement of Dygraph into Static Graph.
    """

    def __init__(self, wrapper_root):
        assert isinstance(
            wrapper_root, AstNodeWrapper
        ), "Type of input node should be AstNodeWrapper, but received %s ." % type(
            wrapper_root)
        self.root = wrapper_root.node
        self.static_analysis_visitor = StaticAnalysisVisitor(self.root)

    def transform(self):
        """
        Main function to transform AST.
        """
        self.visit(self.root)

    def visit_If(self, node):
        self.generic_visit(node)
        new_vars_stmts, true_func_node, false_func_node, return_name_ids = transform_if_else(
            node, self.root)

        new_node = create_convert_ifelse_node(return_name_ids, node.test,
                                              true_func_node, false_func_node)

        return new_vars_stmts + [true_func_node, false_func_node] + [new_node]

    def visit_Call(self, node):
        # Remove `numpy()` statement, like `Tensor.numpy()[i]` -> `Tensor[i]`
        if isinstance(node.func, gast.Attribute):
            attribute = node.func
            if attribute.attr == 'numpy':
                node = attribute.value
        self.generic_visit(node)
        return node

    def visit_IfExp(self, node):
        """
        Transformation with `true_fn(x) if Tensor > 0 else false_fn(x)`
        """
        self.generic_visit(node)

        new_node = create_convert_ifelse_node(None, node.test, node.body,
                                              node.orelse, True)
        # Note: A blank line will be added separately if transform gast.Expr
        # into source code. Using gast.Expr.value instead to avoid syntax error
        # in python.
        if isinstance(new_node, gast.Expr):
            new_node = new_node.value

        return new_node


class NameVisitor(gast.NodeVisitor):
    def __init__(self, end_node=None):
        # The terminate node of the visitor.
        self.end_node = end_node
        # Dict to store the names and ctxs of vars.
        self.name_ids = defaultdict(list)
        # List of current visited nodes
        self.ancestor_nodes = []
        # Available only when end_node is set.
        self._is_finished = False
        self._candidate_ctxs = (gast.Store, gast.Load, gast.Param)
        self._def_func_names = set()

    def visit(self, node):
        """Visit a node."""
        if node == self.end_node or self._is_finished:
            self._is_finished = True
            return

        self.ancestor_nodes.append(node)
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        ret = visitor(node)
        self.ancestor_nodes.pop()

        return ret

    def visit_If(self, node):
        """
        For nested `if/else`, the created vars are not always visible for parent node.
        In addition, the vars created in `if.body` are not visible for `if.orelse`.

        Case 1:
            x = 1
            if m > 1:
                res = new_tensor
            res = res + 1   # Error, `res` is not visible here.

        Case 2:
            if x_tensor > 0:
                res = new_tensor
            else:
                res = res + 1   # Error, `res` is not visible here.

        In above two cases, we should consider to manage the scope of vars to parsing
        the arguments and returned vars correctly.
        """
        if not self.end_node:
            self.generic_visit(node)
        else:
            before_if_name_ids = copy.deepcopy(self.name_ids)
            body_name_ids = self._visit_child(node.body)
            # If traversal process stops early in `if.body`, return the currently seen name_ids.
            if self._is_finished:
                self._update_name_ids(before_if_name_ids)
            else:
                else_name_ids = self._visit_child(node.orelse)
                # If traversal process stops early in `if.orelse`, return the currently seen name_ids.
                if self._is_finished:
                    self._update_name_ids(before_if_name_ids)
                else:
                    # Blocks the vars in `if.body` and only inserts the vars both created in 'if/else' branch
                    # into name_ids.
                    new_name_ids = self._find_new_name_ids(body_name_ids,
                                                           else_name_ids)
                    for new_name_id in new_name_ids:
                        before_if_name_ids[new_name_id].append(gast.Store())

                    self.name_ids = before_if_name_ids

    def visit_Attribute(self, node):
        if not self._is_call_func_name_node(node):
            self.generic_visit(node)

    def visit_Name(self, node):
        blacklist = {'True', 'False', 'None'}
        if node.id in blacklist: return
        if node.id in self._def_func_names:
            return
        if not self._is_call_func_name_node(node):
            if isinstance(node.ctx, self._candidate_ctxs):
                self.name_ids[node.id].append(node.ctx)

    def visit_Assign(self, node):
        # Visit `value` firstly.
        node._fields = ('value', 'targets')
        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        self._def_func_names.add(node.name)
        if not self.end_node:
            self.generic_visit(node)
        else:
            before_name_ids = copy.deepcopy(self.name_ids)
            self.name_ids = defaultdict(list)
            self.generic_visit(node)

            if self._is_finished:
                self._update_name_ids(before_name_ids)
            else:
                self.name_ids = before_name_ids

    def _visit_child(self, node):
        self.name_ids = defaultdict(list)
        if isinstance(node, list):
            for item in node:
                if isinstance(item, gast.AST):
                    self.visit(item)
        elif isinstance(node, gast.AST):
            self.visit(node)

        return copy.deepcopy(self.name_ids)

    def _find_new_name_ids(self, body_name_ids, else_name_ids):
        def is_required_ctx(ctxs, required_ctx):
            for ctx in ctxs:
                if isinstance(ctx, required_ctx):
                    return True
            return False

        candidate_name_ids = set(body_name_ids.keys()) & set(else_name_ids.keys(
        ))
        store_ctx = gast.Store
        new_name_ids = set()
        for name_id in candidate_name_ids:
            if is_required_ctx(body_name_ids[name_id],
                               store_ctx) and is_required_ctx(
                                   else_name_ids[name_id], store_ctx):
                new_name_ids.add(name_id)

        return new_name_ids

    def _is_call_func_name_node(self, node):
        if len(self.ancestor_nodes) > 1:
            assert self.ancestor_nodes[-1] == node
            parent_node = self.ancestor_nodes[-2]
            if isinstance(parent_node, gast.Call) and parent_node.func == node:
                return True
        return False

    def _update_name_ids(self, new_name_ids):
        for name_id, ctxs in six.iteritems(new_name_ids):
            self.name_ids[name_id] = ctxs + self.name_ids[name_id]


def get_name_ids(nodes, end_node=None):
    """
    Return all ast.Name.id of python variable in nodes.
    """
    name_visitor = NameVisitor(end_node)
    for node in nodes:
        name_visitor.visit(node)
    return name_visitor.name_ids


def parse_cond_args(var_ids_dict, return_ids=None, ctx=gast.Load):
    """
    Find out the ast.Name.id list of input by analyzing node's AST information.
    """

    name_ids = [
        var_id for var_id, var_ctx in six.iteritems(var_ids_dict)
        if isinstance(var_ctx[0], ctx)
    ]
    if return_ids:
        new_args = set(return_ids) - set(name_ids)
        name_ids.extend(list(new_args))
    name_ids.sort()
    args = [
        gast.Name(
            id=name_id, ctx=gast.Load(), annotation=None, type_comment=None)
        for name_id in name_ids
    ]
    arguments = gast.arguments(
        args=args,
        posonlyargs=[],
        vararg=None,
        kwonlyargs=[],
        kw_defaults=None,
        kwarg=None,
        defaults=[])

    return arguments


def parse_cond_return(parent_vars_dict, if_vars_dict, else_vars_dict,
                      after_ifelse_vars_dict):
    """
    Find out the ast.Name list of output by analyzing node's AST information.
    One of the following conditions should be satisfied while determining whether a variable is a return value:
    1. the var in parent scope is modified in If.body or If.orelse node.
    2. new var is both created in If.body and If.orelse node.
    3. new var is created only in one of If.body or If.orelse node, and it used as gast.Load firstly after gast.If node.

    For example:
        x, y = 5, 10
        if x > 4:
            x = x+1
            z = x*x
            q = 10
        else:
            y = y - 1
            z = y*y
            m = 20
            n = 20

        print(q)
        n = 30
        print(n)


    The return_ids are (x, y, z, q) for `If.body` and `If.orelse`node, because
    1. x is modified in If.body node,
    2. y is modified in If.body node,
    3. z is both created in If.body and If.orelse node,
    4. q is created only in If.body, and it is used by `print(q)` as gast.Load.
    Note:
        After transformed, q and z are created in parent scope. For example,

        x, y = 5, 10
        q = paddle.jit.dy2static.data_layer_not_check(name='q', shape=[-1], dtype='float32')
        z = paddle.jit.dy2static.data_layer_not_check(name='z', shape=[-1], dtype='float32')

        def true_func(x, y, q):
            x = x+1
            z = x*x
            q = 10
            return x,y,z,q

        def false_func(x, y, q):
            y = y - 1
            z = y*y
            m = 20
            n = 20
            return x,y,z,q

        x,y,z,q = fluid.layers.cond(x>4, lambda: true_func(x, y), lambda: false_func(x, y, q))

    m and n are not in return_ids, because
    5. m is created only in If.orelse, but it is not used after gast.If node.
    6. n is created only in If.orelse, and it is used by `n = 30` and `print(n)`, but it is not used as gast.Load firstly but gast.Store .

    """

    def _is_return_var(ctxs):
        for ctx in ctxs:
            if isinstance(ctx, (gast.Store, gast.Param)):
                return True
        return False

    def _vars_with_store(ids_dict):
        vars = []
        for k, ctxs in six.iteritems(ids_dict):
            if _is_return_var(ctxs):
                vars.append(k)
        return vars

    def _modified_vars(child_dict, parent_dict):
        return set([
            var for var in _vars_with_store(child_dict) if var in parent_dict
        ])

    def _vars_loaded_before_store(ids_dict):
        new_dict = defaultdict(list)
        for k, ctxs in six.iteritems(ids_dict):
            for ctx in ctxs:
                if isinstance(ctx, gast.Load):
                    new_dict[k].append(ctx)
                elif isinstance(ctx, gast.Store):
                    break
        return new_dict

    # modified vars
    body_modified_vars = _modified_vars(if_vars_dict, parent_vars_dict)
    orelse_modified_vars = _modified_vars(else_vars_dict, parent_vars_dict)
    modified_vars = body_modified_vars | orelse_modified_vars

    # new vars
    body_new_vars = set([
        var for var in _vars_with_store(if_vars_dict)
        if var not in parent_vars_dict
    ])
    orelse_new_vars = set([
        var for var in _vars_with_store(else_vars_dict)
        if var not in parent_vars_dict
    ])
    new_vars_in_body_or_orelse = body_new_vars | orelse_new_vars
    new_vars_in_one_of_body_or_orelse = body_new_vars ^ orelse_new_vars

    # 1. the var in parent scope is modified in If.body or If.orelse node.
    modified_vars_from_parent = modified_vars - new_vars_in_body_or_orelse

    # 2. new var is both created in If.body and If.orelse node.
    new_vars_in_body_and_orelse = body_new_vars & orelse_new_vars

    # 3. new var is created only in one of If.body or If.orelse node, and it used as gast.Load firstly after gast.If node.
    used_vars_after_ifelse = set(
        [var for var in _vars_loaded_before_store(after_ifelse_vars_dict)])
    new_vars_to_create = new_vars_in_one_of_body_or_orelse & used_vars_after_ifelse | new_vars_in_body_and_orelse

    # 4. generate return_ids of if/else node.
    return_ids = list(modified_vars_from_parent | new_vars_in_body_and_orelse |
                      new_vars_to_create)
    return_ids.sort()

    return return_ids, modified_vars_from_parent, new_vars_to_create


def transform_if_else(node, root):
    """
    Transform ast.If into control flow statement of Paddle static graph.
    """
    # TODO(liym27): Consider variable like `self.a` modified in if/else node.
    parent_name_ids = get_name_ids([root], end_node=node)
    body_name_ids = get_name_ids(node.body)
    orelse_name_ids = get_name_ids(node.orelse)

    # Get after_ifelse_name_ids, which means used var names after If.body and If.orelse node.
    after_ifelse_name_ids = defaultdict(list)
    all_name_ids = get_name_ids([root])
    for name in all_name_ids:
        before_var_names_ids = parent_name_ids.get(name, []) + \
                           body_name_ids.get(name, []) + orelse_name_ids.get(name, [])
        # Note: context of node.Name like gast.Load is a concrete object which has unique id different from other gast.Load
        #  E.g. ctx of `x` can be [<gast.Load object at 0x142a33c90>, <gast.Load object at 0x142a51950>, <gast.Param object at 0x1407d8250>]
        after_var_names_ids = [
            ctx for ctx in all_name_ids[name] if ctx not in before_var_names_ids
        ]
        if after_var_names_ids:
            after_ifelse_name_ids[name] = after_var_names_ids

    return_name_ids, modified_name_ids_from_parent, new_vars_to_create = parse_cond_return(
        parent_name_ids, body_name_ids, orelse_name_ids, after_ifelse_name_ids)

    # NOTE: Python can create variable only in if body or only in else body, and use it out of if/else.
    # E.g.
    #
    # if x > 5:
    #   a = 10
    # print(a)
    #
    # Create static variable for those variables
    create_new_vars_in_parent_stmts = []
    for name in new_vars_to_create:
        # NOTE: Consider variable like `self.a` modified in if/else node.
        if "." not in name:
            create_new_vars_in_parent_stmts.append(
                create_static_variable_gast_node(name))

    modified_name_ids = modified_name_ids_from_parent | new_vars_to_create

    true_func_node = create_funcDef_node(
        node.body,
        name=unique_name.generate(TRUE_FUNC_PREFIX),
        input_args=parse_cond_args(body_name_ids, modified_name_ids),
        return_name_ids=return_name_ids)
    false_func_node = create_funcDef_node(
        node.orelse,
        name=unique_name.generate(FALSE_FUNC_PREFIX),
        input_args=parse_cond_args(orelse_name_ids, modified_name_ids),
        return_name_ids=return_name_ids)
    return create_new_vars_in_parent_stmts, true_func_node, false_func_node, return_name_ids


def create_convert_ifelse_node(return_name_ids,
                               pred,
                               true_func,
                               false_func,
                               is_if_expr=False):
    """
    Create `paddle.jit.dy2static.convert_ifelse(
            pred, true_fn, false_fn, true_args, false_args, return_vars)`
    to replace original `python if/else` statement.
    """

    def create_name_nodes(name_ids):
        if not name_ids:
            return gast.Tuple(elts=[], ctx=gast.Load())

        gast_names = [
            gast.Name(
                id=name_id, ctx=gast.Load(), annotation=None, type_comment=None)
            for name_id in name_ids
        ]
        name_node = gast.Tuple(elts=gast_names, ctx=gast.Load())
        return name_node

    if is_if_expr:
        true_args = gast.Tuple(elts=[], ctx=gast.Load())
        false_args = gast.Tuple(elts=[], ctx=gast.Load())
        true_func_source = "lambda : {}".format(ast_to_source_code(true_func))
        false_func_source = "lambda : {}".format(ast_to_source_code(false_func))
    else:
        true_args = gast.Tuple(elts=true_func.args.args, ctx=gast.Load())
        false_args = gast.Tuple(elts=false_func.args.args, ctx=gast.Load())
        true_func_source = true_func.name
        false_func_source = false_func.name

    return_vars = create_name_nodes(return_name_ids)

    convert_ifelse_layer = gast.parse(
        'paddle.jit.dy2static.convert_ifelse('
        '{pred}, {true_fn}, {false_fn}, {true_args}, {false_args}, {return_vars})'.
        format(
            pred=ast_to_source_code(pred),
            true_fn=true_func_source,
            false_fn=false_func_source,
            true_args=ast_to_source_code(true_args),
            false_args=ast_to_source_code(false_args),
            return_vars=ast_to_source_code(return_vars))).body[0].value

    if return_name_ids:
        _, cond_node = create_assign_node(return_name_ids, convert_ifelse_layer)
    else:  # No variables can be returned if no assign statement in if.body.
        cond_node = gast.Expr(value=convert_ifelse_layer)

    return cond_node
