# 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.fluid.data_feeder import convert_dtype
from paddle.fluid.dygraph.dygraph_to_static.variable_trans_func import to_static_variable
from paddle.fluid.framework import core, Variable
from paddle.fluid.layers import Assert, Print
from paddle.fluid.layers import array_length, array_read, array_write, create_array
from paddle.fluid.layers import assign, fill_constant, slice, reduce_all, reduce_any
from paddle.fluid.layers import cast, control_flow, logical_and, logical_not, logical_or, nn
from paddle.fluid.layers.control_flow import cond, while_loop, less_than, increment


def convert_while_loop(cond, body, loop_vars):
    """
    A function representation of a Python ``while`` statement.

    Args:
        cond(Callable): A callable object that returns a boolean variable to control whether to execute the loop body. It takes ``loop_vars`` as arguments.
        body(Callable): A callable object that returns a tuple or list of variables with the same arguments ``loops_vars`` as ``cond`` .
        loop_vars(list|tuple): A list or tuple of variables passed to ``cond`` and ``body`` .

    Returns:
        A list or tuple of variables which returned by ``body``.
    """

    # NOTE: It may be slower if cond is very expensive, but usually cond is just O(1).
    # If loop_vars is changed during cond callable, then it causes bug, but current logical_and/logical_not/... doesn't change the loop_vars.
    pred = cond(*loop_vars)
    if isinstance(pred, Variable):
        loop_vars = _run_paddle_while_loop(cond, body, loop_vars)
    else:
        loop_vars = _run_py_while(cond, body, loop_vars)

    return loop_vars


def _run_paddle_while_loop(cond, body, loop_vars):
    # NOTE: loop_vars of Paddle op `control_flow.while_loop` must be Paddle Tensors.
    loop_vars = [to_static_variable(var) for var in loop_vars]
    loop_vars = control_flow.while_loop(cond, body, loop_vars)
    return loop_vars


def _run_py_while(cond, body, loop_vars):
    while cond(*loop_vars):
        loop_vars = body(*loop_vars)
    return loop_vars


def convert_logical_and(x_func, y_func):
    """
    A function representation of a Python ``and`` statement.

    Args:
        x_func(callable): x_func() is the left hand operand of ``and`` operator. x_func() is bool or Tensor.
        y_func(callable): y_func() is the right hand operand of ``and`` operator.  y_func() is bool or Tensor.

    Returns:
        A python bool variable or a bool Tensor.

    NOTE(liym27):
        1) The operands are executed sequentially according to the running logic of Python. So here the arguments
        should be callable.
        2) If the left hand operand is False, the right hand operand should be executed.

        For example:
            a = x > 1 and y < 1
        Transformed code:
            a = paddle.jit.dy2static.convert_logical_and(lambda:x>1, lambda:y<1)

          In `convert_logical_and(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`. And
        if `x>1` is False, `y<1` should NOT be run.
    """
    x_value = x_func()
    if not isinstance(x_value, Variable):
        return _run_py_logical_and(lambda: x_value, y_func)

    y_value = y_func()
    if not isinstance(y_value, Variable):
        return _run_py_logical_and(lambda: y_value, lambda: x_value)

    return _run_paddle_logical_and(x_value, y_value)


def _run_paddle_logical_and(x, y):
    x = cast_bool_if_necessary(x)
    y = cast_bool_if_necessary(y)
    return logical_and(x, y)


def _run_py_logical_and(x_func, y_func):
    x_value = x_func()
    assert not isinstance(x_value, Variable)

    # NOTE(liym27):
    #  1. Returns y_func() if x_value is False;
    #  2. If x_value is False, y_func() should not be run.
    return x_value and y_func()


def convert_logical_or(x_func, y_func):
    """
    A function representation of a Python ``or`` statement.

    Args:
        x_func(callable): x_func() is the left hand operand of ``or`` operator. x_func() is bool or Tensor.
        y_func(callable): y_func() is the right hand operand of ``or`` operator.  y_func() is bool or Tensor.

    Returns:
        A python bool variable or a bool Tensor.

    NOTE(liym27):
        1) The operands are executed sequentially according to the running logic of Python. So here the arguments
        should be callable.
        2) If the left hand operand is True, the right hand operand should be executed.

        For example:
            a = x > 1 or y < 1
        Transformed code:
            a = paddle.jit.dy2static.convert_logical_or(lambda:x>1, lambda:y<1)

        In `convert_logical_or(lambda:x>1, lambda:y<1)`, `lambda:y<1` must be run after `lambda:x>1`. And
        if `x>1` is True, `y<1` should NOT be run.
    """
    x_value = x_func()
    if not isinstance(x_value, Variable):
        return _run_py_logical_or(lambda: x_value, y_func)

    y_value = y_func()
    if not isinstance(y_value, Variable):
        return _run_py_logical_or(lambda: y_value, lambda: x_value)

    return _run_paddle_logical_or(x_value, y_value)


def _run_paddle_logical_or(x, y):
    x = cast_bool_if_necessary(x)
    y = cast_bool_if_necessary(y)
    return logical_or(x, y)


def _run_py_logical_or(x_func, y_func):
    x_value = x_func()
    assert not isinstance(x_value, Variable)

    # NOTE(liym27):
    #  1. Returns y_func() if x_value is False;
    #  2. If x_value is True, y_func() should not be run.
    return x_value or y_func()


def convert_logical_not(x):
    """
    A function representation of a Python ``not`` statement.

    Args:
        x(bool|Tensor): Operand of of ``not`` operator.

    Returns:
        A python bool variable or a bool Tensor.
    """

    if isinstance(x, Variable):
        return _run_paddle_logical_not(x)
    else:
        return _run_py_logical_not(x)


def _run_paddle_logical_not(x):
    x = cast_bool_if_necessary(x)
    return logical_not(x)


def _run_py_logical_not(x):
    return not x


def convert_ifelse(pred, true_fn, false_fn, true_args, false_args, return_vars):
    """
    A function representation of a Python ``if/else`` statement.

    Args:
        pred(bool|Tensor): A boolean Tensor which determines whether to return the result of ``true_fn`` or ``false_fn`` .
        true_fn(callable): A callable to be performed if ``pred`` is true.
        false_fn(callable): A callable to be performed if ``pred`` is false.
        true_args(tuple): Parameters of ``true_fn``.
        false_args(tuple): Parameters of ``false_fn``.
        return_vars(tuple): Return variables of ``true_fn`` and ``false_fn``.

    Returns:
        ``true_fn(true_args)`` if the predicate ``pred`` is true else ``false_fn(false_args)`` .

    """
    if isinstance(pred, Variable):
        return _run_paddle_cond(pred, true_fn, false_fn, true_args, false_args,
                                return_vars)
    else:
        return _run_py_ifelse(pred, true_fn, false_fn, true_args, false_args)


def _run_paddle_cond(pred, true_fn, false_fn, true_args, false_args,
                     return_vars):

    return_var_ids = [id(var) for var in return_vars]
    # NOTE 1: Returned vars of Paddle op `control_flow.cond` must be Paddle Tensors
    # NOTE 2: Here uses id(var) not var, because `if var in return_var` use operator `==`,
    #  which will call `fluid.layers.equal` and causes error when var in return_vars is not initialized.
    true_args = [
        to_static_variable(var) if id(var) in return_var_ids else var
        for var in true_args
    ]
    false_args = [
        to_static_variable(var) if id(var) in return_var_ids else var
        for var in false_args
    ]

    pred = cast_bool_if_necessary(pred)
    return control_flow.cond(pred, lambda: true_fn(*true_args),
                             lambda: false_fn(*false_args))


def _run_py_ifelse(pred, true_fn, false_fn, true_args, false_args):
    return true_fn(*true_args) if pred else false_fn(*false_args)


def convert_len(var):
    """
    Returns variable(length) from shape ops based on var.type

    Note: In addition to some ast transformations, some block-related
          operations are added in `len` transformation, such as appending
          `shape_op` in var.block.
    """
    if isinstance(var, Variable):
        if var.type in [
                core.VarDesc.VarType.LOD_TENSOR,
                core.VarDesc.VarType.SELECTED_ROWS
        ]:
            # Note: Length of var may be known ahead of time in dygraph,
            # but it probably represents batch size which can be variant.
            # so we return a variable dynamically inferred from var.shape.
            return nn.shape(var)[0]
        elif var.type == core.VarDesc.VarType.LOD_TENSOR_ARRAY:
            return control_flow.array_length(var)
        else:
            raise TypeError(
                'len(var) only supports LoDTensor/LoDTensorArray/SelectedRows, but received %s.'
                % type(var))
    else:
        return len(var)


def convert_var_shape(x, idx=None, in_control_flow=False):
    """
    A function representation of the shape of variable.
    """

    def has_negative(list_shape, idx=None):
        if idx is not None:
            return list_shape[idx] < 0

        num_negative = sum([1 if i < 0 else 0 for i in list_shape])
        return num_negative > 0

    # When `x` is Variable, call nn.shape(x) in following cases:
    #  (1) The shape of `x` is used in control flow condition.
    #      ```
    #      if x.shape[0] == 1:
    #          y = XX
    #      ```
    #  (2) The dim to be used is negative
    #      ```
    #      # Assume x.shape=[3, -1] in static mode
    #      y = paddle.reshape(x, shape=[1, x.shape[1]])
    #      ```
    if isinstance(x, Variable) and (in_control_flow or has_negative(x.shape,
                                                                    idx)):
        return nn.shape(x) if idx is None else nn.shape(x)[idx]
    else:
        return x.shape if idx is None else x.shape[idx]


def convert_var_shape_simple(x):
    """
    A function representation of the shape of variable.
    """
    if isinstance(x, Variable):
        return nn.shape(x)
    else:
        return x.shape


def eval_if_exist_else_none(name, global_symbol_table):
    """
    Args:
        name([str]): Expression passed into `eval`.
        local_symbol_table(dict): Specified from `globals()`. DO NOT use `locals()`,
                                  because all STATIC_CONVERT_VAR_SHAPE_SUFFIX vars is
                                  declared with keyword `global`.
    
    Returns:
        Return the variable if found in global_symbol_table else None.
    """
    try:
        return eval(name, global_symbol_table)
    except:
        return None


def choose_shape_attr_or_api(attr_shape, api_shape, idx=None):
    """
    Input can be attribute `x.shape` or api `shape(x)`, this function
    chooses which one to return to use in dy2stat.

    Note: sometimes users write `x.shape[3]`, so attr_shape can be an integer.
    """
    if api_shape is None:
        return attr_shape if idx is None else attr_shape[idx]
    if not isinstance(attr_shape, (list, tuple)):
        # some variables like x.shape[0] is no longer a list or tuple
        if isinstance(attr_shape, int) and attr_shape < 0:
            return api_shape if idx is None else api_shape[idx]
        return attr_shape if idx is None else attr_shape[idx]

    def has_negative(list_shape, idx=None):
        if idx is not None:
            return list_shape[idx] < 0

        num_negative = sum([1 if i < 0 else 0 for i in list_shape])
        return num_negative > 0

    if has_negative(attr_shape, idx):
        return api_shape if idx is None else api_shape[idx]
    return attr_shape if idx is None else attr_shape[idx]


def convert_shape_compare(left, *args):
    """
    A function handles comparison difference between Paddle and Python.
    For example, if x and y are Tensors, x.shape == y.shape will return single
    boolean Value (True/False). However, paddle.shape(x) == paddle.shape(y) is
    an element-wise comparison. The difference can cause dy2stat error. So we
    create this function to handle the difference.

    Args:
        left: variable
        *args: compare_op(str), variable, compare_op(str), variable, where
            compare_op means "<", ">", "==", "!=", etc.
    Returns:
        If the variables to compare are NOT Paddle Variables, we will return as
        Python like "a op1 b and b op2 c and ... ".
        If the variables to compare are Paddle Variables, we will do elementwise
        comparsion first and then reduce to a boolean whose numel is 1.
        
    """
    args_len = len(args)
    assert args_len >= 2, "convert_shape_compare needs at least one right compare variable"
    assert args_len % 2 == 0, "Illegal input for convert_shape_compare, *args should be op(str), var, op(str), var ..."
    num_cmp = args_len // 2
    if isinstance(left, Variable):

        def reduce_compare(x, op_str, y):
            element_wise_result = eval("x " + op_str + " y")
            if op_str == "!=":
                return reduce_any(element_wise_result)
            elif op_str == "is" or op_str == "is not" or op_str == "in" or op_str == "not in":
                return element_wise_result
            else:
                return reduce_all(element_wise_result)

        final_result = reduce_compare(left, args[0], args[1])
        for i in range(1, num_cmp):
            cmp_left = args[i * 2 - 1]
            cmp_op = args[i * 2]
            cmp_right = args[i * 2 + 1]
            cur_result = reduce_compare(cmp_left, cmp_op, cmp_right)
            final_result = convert_logical_and(lambda: final_result,
                                               lambda: cur_result)
        return final_result
    else:
        cmp_left = left
        final_result = None
        for i in range(num_cmp):
            cmp_op = args[i * 2]
            cmp_right = args[i * 2 + 1]
            cur_result = eval("cmp_left " + cmp_op + " cmp_right")
            if final_result is None:
                final_result = cur_result
            else:
                final_result = final_result and cur_result

            if final_result is False:
                return False
            cmp_left = cmp_right
        return final_result


def cast_bool_if_necessary(var):
    assert isinstance(var, Variable)
    if convert_dtype(var.dtype) not in ['bool']:
        var = cast(var, dtype="bool")
    return var


def convert_var_dtype(var, dtype):
    if isinstance(var, Variable):
        src_dtype = convert_dtype(var.dtype)
        assert src_dtype in [
            'bool', 'float16', 'float32', 'float64', 'int32', 'int64', 'uint8'
        ], "The dtype of var {} is {}, which is not supported in the cast op.".format(
            var.name, src_dtype)
        assert dtype in [
            'bool', 'int', 'float'
        ], "The casted target dtype is {}, which is not supported in type casting.".format(
            dtype)
        cast_map = {
            'bool': 'bool',
            'int': 'int32',
            'float': 'float32',
        }
        return cast(var, dtype=cast_map[dtype])
    else:
        return eval('{}(var)'.format(dtype))


def convert_assert(cond, message=""):
    """
    A function representation of a Python ``assert`` statement.
    """
    if isinstance(cond, Variable):
        cond = cast(cond, "bool")
        # NOTE: message is not used because Paddle Assert has no corresponding parameter to use.
        return Assert(cond)
    else:
        assert cond, message


def convert_print(*args):
    """
    A function representing Python ``print`` statement. Note: this is a basic
    python function so we haven't handle sep, end, file and flush parameters of
    python function.
    """
    for var in args:
        if isinstance(var, Variable):
            var = Print(var)
        else:
            print(var)


def convert_pop(target, *args):
    """
    A function representation of a Python pop statement for a list or dict.

    Args:
        target(list|dict|Tensor): A variable to pop item from.
        *args(tuple): index or default value to parse.

    Returns:
        A item poped from target.
    """

    is_variable = isinstance(target, Variable)
    if is_variable:
        is_tensor_array = target.type == core.VarDesc.VarType.LOD_TENSOR_ARRAY

    if is_variable and is_tensor_array:
        return _run_paddle_pop(target, *args)
    else:
        return _run_python_pop(target, *args)


def _run_paddle_pop(array, *args):
    if len(args) == 0:
        idx = -1
    else:
        idx = args[0]

    assert isinstance(idx, int)

    def cond(i, new_array):
        return less_than(i, arr_len)

    def body(i, new_array):
        item = array_read(array=array, i=i)
        array_write(item, array_length(new_array), new_array)
        i = increment(i)
        return i, new_array

    arr_len = array_length(array)
    if idx < 0:
        idx = idx + arr_len
    else:
        idx = fill_constant(shape=[1], dtype="int64", value=idx)

    pop_item = array_read(array, idx)

    new_array = _slice_tensor_array(array, 0, idx)
    i = idx + 1
    _, new_array = while_loop(cond, body, [i, new_array])
    assign(input=new_array, output=array)

    return pop_item


# TODO(liym27): A better way to slice tensor array.
#  Maybe support start == end for slice op.
def _slice_tensor_array(array, start, end):
    def true_fn():
        null_array = create_array("float32")
        return null_array

    def false_fn(array, start, end):
        new_array = slice(array, starts=[start], ends=[end], axes=[0])
        return new_array

    new_array = cond(start == end, true_fn, lambda: false_fn(array, start, end))
    return new_array


def _run_python_pop(target, *args):
    # 1. pop for a dict
    if len(args) == 2:
        idx, default = args
        return target.pop(idx, default)

    # 2. pop for a list or dict
    else:
        idx = args[0] if args else -1
        return target.pop(idx)
