# Copyright (c) 2022 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 paddle
import numpy as np
from paddle import _legacy_C_ops
from paddle.fluid.framework import (
    _test_eager_guard,
    Variable,
    _in_legacy_dygraph,
)
from paddle.fluid import core
from paddle.fluid.layers.utils import _hash_with_id
from paddle.fluid.dygraph.base import switch_to_static_graph
from paddle.fluid.executor import (
    _is_enable_standalone_executor,
    _is_dy2st_enable_standalone_executor,
)

import unittest


def _append_backward_desc(main_program, outs):
    # make sure all status of is_test are False in train mode.
    program = main_program.clone()
    targets = []
    for out in outs:
        if isinstance(out, Variable):
            targets.append(program.global_block().var(out.name))

    if targets:
        paddle.fluid.backward.gradients(targets=targets, inputs=[])

    return program


# def _set_grad_type(params, train_program):
#     # NOTE: if user set sparse gradient mode, the param's gradient
#     # will be SelectedRows, not LoDTensor. But tracer will just
#     # set param grad VarBase by forward VarBase(LoDTensor)
#     # If we don't change grad_var type here, RunProgramOp need
#     # transform SelectedRows to LoDTensor forcibly, it may not
#     # be user wanted result.
#     for param in params:
#         grad_name = param.name + core.grad_var_suffix()
#         grad_var = train_program.desc.block(0).find_var(
#             grad_name.encode())
#         # NOTE: cannot find var desc maybe no problem, such as in batch_norm
#         if grad_var is None:
#             continue
#         param._set_grad_type(grad_var.type())


def _create_out(var):
    assert isinstance(var, Variable)
    var_desc = var.desc
    varbase = None
    if _in_legacy_dygraph():
        var_base = core.VarBase(
            var_desc.dtype(),
            var_desc.shape(),
            var_desc.name(),
            var_desc.type(),
            False,
        )
    else:
        var_base = core.eager.Tensor(
            var_desc.dtype(),
            var_desc.shape(),
            var_desc.name(),
            var_desc.type(),
            False,
        )
    return var_base


@switch_to_static_graph
def _add_build_strategy_for(input_program, start_op_index, end_op_index):
    compiled_program = paddle.static.CompiledProgram(
        core.Graph(input_program.desc, start_op_index, end_op_index),
        build_strategy=paddle.static.BuildStrategy(),
    )
    compiled_program._compile(
        core.Scope(), paddle.framework._current_expected_place()
    )
    ir_graph = paddle.fluid.framework.IrGraph(compiled_program._graph)
    builded_program = ir_graph.to_program()
    return builded_program


class TestRunProgram(unittest.TestCase):
    def test_eager(self):
        paddle.set_device('cpu')
        paddle.enable_static()
        # step 1: construct program
        x = paddle.static.data(shape=[2, 4], name='x')
        x.stop_gradient = False
        y = paddle.static.data(shape=[4, 2], name='y')
        y.stop_gradient = False
        out = paddle.matmul(x, y)

        main_program = paddle.static.default_main_program()
        program = _append_backward_desc(main_program, [out])
        forward_program = _add_build_strategy_for(
            program, 0, main_program.desc.block(0).op_size()
        )
        backward_program = _add_build_strategy_for(
            program,
            main_program.desc.block(0).op_size() + 2,
            program.desc.block(0).op_size(),
        )

        paddle.disable_static('cpu')
        # step 2: call run_program in eager mode
        with _test_eager_guard():
            x_t = paddle.ones([2, 4])
            x_t.name = "x"
            x_t.stop_gradient = False
            y_t = paddle.ones([4, 2])
            y_t.name = "y"
            y_t.stop_gradient = False

            fake_var = paddle.zeros([1])
            fake_var.name = 'Fake_var'

            out_t = _create_out(out)

            scope = core.Scope()
            attrs = [
                'global_block',
                program.desc.block(0),
                'start_op_index',
                0,
                'end_op_index',
                main_program.desc.block(0).op_size(),
                'is_test',
                False,
                'program_id',
                _hash_with_id(program),
            ]

            use_interpretorcore = (
                _is_enable_standalone_executor()
                and _is_dy2st_enable_standalone_executor()
            )
            attrs.extend(('use_interpretorcore', use_interpretorcore))
            if use_interpretorcore:
                attrs.extend(
                    (
                        'forward_global_block',
                        forward_program.desc.block(0),
                        'backward_global_block',
                        backward_program.desc.block(0),
                    )
                )

            _legacy_C_ops.run_program(
                [x_t, y_t],
                [fake_var],
                [out_t],
                [scope],
                [fake_var],
                None,
                *attrs
            )

            loss = paddle.mean(out_t)
            loss.backward()

            np.testing.assert_array_equal(np.ones([2, 2]) * 4, out_t.numpy())
            np.testing.assert_array_equal(
                np.ones([2, 4]) * 0.5, x_t.grad.numpy()
            )
            np.testing.assert_array_equal(
                np.ones([4, 2]) * 0.5, y_t.grad.numpy()
            )


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