# 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 unittest

import paddle
import numpy as np
import paddle.fluid as fluid
from paddle.fluid.dygraph.jit import declarative
from paddle.fluid.layers.utils import map_structure

SEED = 2020
np.random.seed(SEED)


# Situation 1: Test list append
def test_list_append_without_control_flow(x):
    # Python list will not be transformed.
    x = fluid.dygraph.to_variable(x)
    a = []
    # It's a plain python control flow which won't be transformed
    if 2 > 1:
        a.append(x)
    return a


def test_list_append_in_if(x):
    x = fluid.dygraph.to_variable(x)
    a = []
    if x.numpy()[0] > 0:
        a.append(x)
    else:
        a.append(
            fluid.layers.fill_constant(
                shape=[1, 2], value=9, dtype="int64"))
    # TODO(Aurelius84): Currently, run_program_op doesn't support output LoDTensorArray.
    return a[0]


def test_list_append_in_for_loop(x, iter_num):
    x = fluid.dygraph.to_variable(x)
    # Use `fill_constant` so that static analysis can analyze the type of iter_num is Tensor
    iter_num = fluid.layers.fill_constant(
        shape=[1], value=iter_num, dtype="int32"
    )  # TODO(liym27): Delete it if the type of parameter iter_num can be resolved
    a = []
    for i in range(iter_num):
        a.append(x)
    return a[0]


def test_list_append_in_for_subscript(x):
    x = fluid.dygraph.to_variable(x)
    iter_num = paddle.shape(x)[0]
    a = []
    for i in range(iter_num):
        x = x + 1
        a.append(x)
    out = paddle.concat(a)
    return out[0]


def test_list_append_in_while_loop_subscript(x):
    x = fluid.dygraph.to_variable(x)
    iter_num = paddle.shape(x)[0]
    a = []
    i = 0
    while i < iter_num:
        x = x + 1
        a.append(x)
        i += 1
    out = paddle.concat(a)
    return out[0]


def test_list_append_in_for_loop_with_concat(x, iter_num):
    x = fluid.dygraph.to_variable(x)
    a = []
    # Use `fill_constant` so that static analysis can analyze the type of iter_num is Tensor
    iter_num = fluid.layers.fill_constant(
        shape=[1], value=iter_num, dtype="int32"
    )  # TODO(liym27): Delete it if the type of parameter iter_num can be resolved
    for i in range(iter_num):
        a.append(x)
    a = fluid.layers.concat(a, axis=0)
    return a


def test_list_append_in_while_loop(x, iter_num):
    x = fluid.dygraph.to_variable(x)
    iter_num = fluid.layers.fill_constant(
        shape=[1], value=iter_num, dtype="int32")
    a = []
    i = 0
    while i < iter_num:
        a.append(x)
        i += 1
    return a[0]


def test_list_append_in_while_loop_with_stack(x, iter_num):
    x = fluid.dygraph.to_variable(x)
    iter_num = fluid.layers.fill_constant(
        shape=[1], value=iter_num, dtype="int32")
    a = []
    i = 0
    while i < iter_num.numpy()[0]:
        a.append(x)
        i += 1
    out = fluid.layers.stack(a, axis=1)
    return out


# Situation 2: Test list pop
def test_list_pop_without_control_flow_1(x):
    x = fluid.dygraph.to_variable(x)
    a = []
    if 2 > 1:
        a.append(x)
    a.pop()
    return a


def test_list_pop_without_control_flow_2(x):
    x = fluid.dygraph.to_variable(x)
    a = []
    if 2 > 1:
        a.append(x)
        a.append(x + 1)
    last_item = a.pop(1)
    return last_item


def test_list_pop_in_if(x):