# Copyright (c) 2023 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.

import unittest

import numpy as np
from test_case_base import TestCaseBase

import paddle
from paddle.jit.sot.utils.paddle_api_config import add_break_graph_function


def ifelse_func(x, y):
    if x > 0:
        y = y + 1
    else:
        y = y + 2
    return y


class TestIfElse(TestCaseBase):
    def test_simple(self):
        x = paddle.to_tensor([1.0])
        y = paddle.to_tensor([2.0])
        self.assert_results(ifelse_func, x, y)


def multi_output(x: paddle.Tensor):
    m = x + 1
    if x > 0:
        return m
    else:
        return 2 * m


class TestBreakgraph(TestCaseBase):
    def test_simple(self):
        x = paddle.to_tensor(2)
        self.assert_results(multi_output, x)
        x = paddle.to_tensor(-2)
        self.assert_results(multi_output, x)


def print_break_graph(x, y):
    z = x + y
    print(x, z)
    out = y * z * 2
    return out


class TestPrint(TestCaseBase):
    def test_simple(self):
        x = paddle.to_tensor(2)
        y = paddle.to_tensor(3)
        self.assert_results(print_break_graph, x, y)


def to_tensor_break_graph(x, y):
    z = x + y
    out = y * paddle.to_tensor(2) * z
    return out


class TestToTensor(TestCaseBase):
    def test_simple(self):
        add_break_graph_function(paddle.to_tensor)
        x = paddle.to_tensor(2)
        y = paddle.to_tensor(3)
        self.assert_results(to_tensor_break_graph, x, y)


def tensor_clear_gradient(x):
    x = paddle.to_tensor(x)
    x.clear_gradient()
    return x


class TestBreakGraphInResumeFn(TestCaseBase):
    def test_simple(self):
        x = paddle.to_tensor(2)
        self.assert_results(tensor_clear_gradient, x)


def inner_fn(a, b, c, d):
    return a + b * c - d


def multi_stack_args(a, b, c):
    out = inner_fn(a, b, c, paddle.to_tensor(4))
    return out


class TestMultiStackArgs(TestCaseBase):
    def test_simple(self):
        a = paddle.to_tensor(1)
        b = paddle.to_tensor(2)
        c = paddle.to_tensor(3)
        self.assert_results(multi_stack_args, a, b, c)


def break_graph_in_call_method(x):
    out = paddle.nn.functional.relu(paddle.to_tensor([4.0]))
    return x + out


def numpy_break_graph():
    a = paddle.to_tensor([1, 2])
    b = np.sum(a.numpy())
    print(b)
    return b


class TestBreakGraphInCallMethod(TestCaseBase):
    def test_simple(self):
        x = paddle.to_tensor([1.0])
        break_graph_in_call_method(x)
        x = paddle.to_tensor([2.0])
        break_graph_in_call_method(x)

        x = paddle.to_tensor([3.0])
        self.assert_results(break_graph_in_call_method, x)

    def test_numpy(self):
        self.assert_results(numpy_break_graph)


def test_break_graph_repeat(x):
    out = paddle.to_tensor(
        paddle.to_tensor(paddle.to_tensor(paddle.to_tensor([1.0])))
    )
    return x + out


class TestBreakGraphRepeat(TestCaseBase):
    def test_simple(self):
        x = paddle.to_tensor([1.0])
        test_break_graph_repeat(x)
        x = paddle.to_tensor([2.0])
        test_break_graph_repeat(x)

        x = paddle.to_tensor([3.0])
        self.assert_results(test_break_graph_repeat, x)


def break_graph_resume_pass_null(x, y):
    return paddle.add(x, y[0:50] if y is not None else None)


class TestBreakGraphResumePassNull(TestCaseBase):
    def test_break_graph_resume_pass_null(self):
        x = paddle.rand([50, 50], dtype=paddle.float32)
        y = paddle.rand([100, 50], dtype=paddle.float32)
        self.assert_results(break_graph_resume_pass_null, x, y)


class MyLayer(paddle.nn.Layer):
    def __init__(self):
        super().__init__()
        self.head = paddle.nn.Linear(3, 10)

    def forward_features(self, x):
        paddle.jit.sot.psdb.breakgraph()
        return x

    def forward(self, x):
        x = self.forward_features(x)
        return self.head(x)


class TestBreakGraphInLayer(TestCaseBase):
    def test_break_graph_in_layer(self):
        x = paddle.rand([2, 3], dtype=paddle.float32)
        net = MyLayer()
        self.assert_results(net.forward, x)


def dummy(*args):
    return None


def break_graph_call_generator_function(x):
    return dummy(y for y in x)


class TestBreakGraphCallGeneratorFunction(TestCaseBase):
    def test_break_graph_when_call_generator_function(self):
        x = paddle.rand([1], dtype=paddle.float32)
        y = paddle.rand([1], dtype=paddle.float32)
        self.assert_results(break_graph_call_generator_function, [x, y])


def unary_not_break_graph(x):
    return not x


class TestUnaryNot(TestCaseBase):
    def test_unary_not_break_graph(self):
        x = paddle.to_tensor(0)
        self.assert_results(unary_not_break_graph, x)


if __name__ == "__main__":
    unittest.main()