# 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

import paddle
import paddle.distributed as dist


# For API generation which have different type of DistTensor Input and Output
class TestDygraphAPIForDistTensorBranch(unittest.TestCase):
    def check_tensor_eq(self, a, b):
        np1 = a.numpy()
        np2 = b.numpy()
        np.testing.assert_allclose(np1, np2, rtol=1e-05)

    def create_local_and_dist_tensor_pair(self, np_array):
        if np_array.dtype == np.float32:
            local_t = paddle.to_tensor(np_array, dtype='float32')
        elif np_array.dtype == np.float16:
            local_t = paddle.to_tensor(np_array, dtype='float16')
        elif np_array.dtype == np.int32:
            local_t = paddle.to_tensor(np_array, dtype='int32')
        elif np_array.dtype == np.float64:
            local_t = paddle.to_tensor(np_array, dtype='float64')
        elif np_array.dtype == np.bool_:
            local_t = paddle.to_tensor(np_array, dtype='bool')

        mesh = dist.ProcessMesh([0], dim_names=["x"])
        dist_t = dist.shard_tensor(np_array, mesh, [dist.Replicate()])

        local_t.stop_gradient = False
        dist_t.stop_gradient = False

        return local_t, dist_t

    def create_local_and_dist_tensor_list_pair(self, np_array_list):
        assert isinstance(np_array_list, list), (
            "input should be list of np_array!"
        )
        local_t_list = []
        dist_t_list = []
        for np_array in np_array_list:
            local_t, dist_t = self.create_local_and_dist_tensor_pair(np_array)
            local_t_list.append(local_t)
            dist_t_list.append(dist_t)
        return local_t_list, dist_t_list

    def create_two_local_tensor_pair(self, np_array):
        if np_array.dtype == np.float32:
            local_t_1 = paddle.to_tensor(np_array, dtype='float32')
            local_t_2 = paddle.to_tensor(np_array, dtype='float32')
        elif np_array.dtype == np.float16:
            local_t_1 = paddle.to_tensor(np_array, dtype='float16')
            local_t_2 = paddle.to_tensor(np_array, dtype='float16')
        elif np_array.dtype == np.int32:
            local_t_1 = paddle.to_tensor(np_array, dtype='int32')
            local_t_2 = paddle.to_tensor(np_array, dtype='int32')
        elif np_array.dtype == np.bool_:
            local_t_1 = paddle.to_tensor(np_array, dtype='bool')
            local_t_2 = paddle.to_tensor(np_array, dtype='bool')

        local_t_1.stop_gradient = False
        local_t_2.stop_gradient = False

        return local_t_1, local_t_2

    # mixed type of inputs: DenseTensor and DistTensor
    def test_matmul_api_for_mixed_inputs_type(self):
        x = np.random.random(size=[4, 4]).astype("float32")
        y = np.random.random(size=[4, 4]).astype("float32")
        local_x, dist_x = self.create_local_and_dist_tensor_pair(x)
        local_y_1, local_y_2 = self.create_two_local_tensor_pair(y)
        local_out = paddle.matmul(local_x, local_y_1)
        dist_out = paddle.matmul(dist_x, local_y_2)
        self.check_tensor_eq(local_out, dist_out)

        # test backward
        local_out.backward()
        dist_out.backward()
        self.check_tensor_eq(local_x.grad, dist_x.grad)
        self.check_tensor_eq(local_y_1.grad, local_y_2.grad)

    # input: std::vector<phi::Tensor>
    # output: phi::Tensor
    def test_concat_for_dist_tensor(self):
        x1 = np.random.random(size=[4, 4]).astype("float32")
        x2 = np.random.random(size=[4, 4]).astype("float32")
        x3 = np.random.random(size=[4, 4]).astype("float32")
        local_in1, dist_in1 = self.create_local_and_dist_tensor_pair(x1)
        local_in2, dist_in2 = self.create_local_and_dist_tensor_pair(x2)
        local_in3, dist_in3 = self.create_local_and_dist_tensor_pair(x3)
        local_out = paddle.concat([local_in1, local_in2, local_in3])
        dist_out = paddle.concat([dist_in1, dist_in2, dist_in3])
        self.check_tensor_eq(local_out, dist_out)
        local_out.backward()
        dist_out.backward()
        self.check_tensor_eq(local_in1.grad, dist_in1.grad)
        self.check_tensor_eq(local_in2.grad, dist_in2.grad)
        self.check_tensor_eq(local_in3.grad, dist_in3.grad)

    # TODO(GhostScreaming): Support paddle.concat backward later.
    # input: std::vector<phi::Tensor>
    # output: std::vector<phi::Tensor>
    def test_broadcast_tensors_for_dist_tensor(self):
        x1 = np.random.random(size=[4, 4]).astype("float32")
        x2 = np.random.random(size=[4, 4]).astype("float32")
        local_in1, dist_in1 = self.create_local_and_dist_tensor_pair(x1)
        local_in2, dist_in2 = self.create_local_and_dist_tensor_pair(x2)

        local_out1, local_out2 = paddle.broadcast_tensors(
            [local_in1, local_in2]
        )
        dist_out1, dist_out2 = paddle.broadcast_tensors([dist_in1, dist_in2])
        self.check_tensor_eq(local_out1, dist_out1)
        self.check_tensor_eq(local_out2, dist_out2)

        local_out = paddle.concat([local_out1, local_out2])
        dist_out = paddle.concat([dist_out1, dist_out2])

        local_out.backward()
        dist_out.backward()
        self.check_tensor_eq(local_in1.grad, dist_in1.grad)
        self.check_tensor_eq(local_in2.grad, dist_in2.grad)

    # input: paddle::optional<phi::Tensor>
    # output: phi::Tensor
    def test_bincount_api_for_dist_tensor(self):
        x = np.random.random(size=[16]).astype("int32")
        weight = np.random.random(size=[16]).astype("float32")
        local_x, dist_x = self.create_local_and_dist_tensor_pair(x)
        local_weight, dist_weight = self.create_local_and_dist_tensor_pair(
            weight
        )

        local_out = paddle.bincount(local_x, weights=local_weight)
        dist_out = paddle.bincount(dist_x, weights=dist_weight)

        self.check_tensor_eq(local_out, dist_out)

    # input: paddle::optional<std::vector<phi::Tensor>>
    # output: phi::Tensor
    def test_linear_interp_for_dist_tensor(self):
        out_size = np.array(
            [
                50,
            ]
        ).astype("int32")
        shape = [1, 3, 100]
        size1 = np.array([50]).astype("int32")
        scale = 0.5
        scale_list = []
        for _ in range(len(shape) - 2):
            scale_list.append(scale)
        scale = list(map(float, scale_list))

        x = np.random.random(size=shape).astype("float32")
        local_x, dist_x = self.create_local_and_dist_tensor_pair(x)
        local_out_size, dist_out_size = self.create_local_and_dist_tensor_pair(
            out_size
        )
        local_size1, dist_size1 = self.create_local_and_dist_tensor_pair(size1)

        local_scale, dist_scale = self.create_local_and_dist_tensor_pair(
            np.array([0.5]).astype("float32")
        )
        local_out = paddle._C_ops.linear_interp(
            local_x,
            local_out_size,  # Outsize
            [local_size1],  # SizeTensor
            local_scale,  # Scale
            'NCHW',  # data_layout
            -1,  # out_d
            -1,  # out_h
            50,  # in_w * out_w
            scale,
            'linear',  # interp_method
            False,  # align_corners
            1,  # align_mode
        )
        dist_out = paddle._C_ops.linear_interp(
            dist_x,
            dist_out_size,  # Outsize
            [dist_size1],  # SizeTensor
            dist_scale,  # Scale
            'NCHW',  # data_layout
            -1,  # out_d
            -1,  # out_h
            50,  # in_w * out_w
            scale,
            'linear',
            False,  # align_corners
            1,  # align_mode
        )
        self.check_tensor_eq(local_out, dist_out)

    # input: std::vector<phi::Tensor>, phi::Tensor
    # output: inplace std::vector<phi::Tensor>, inplace phi::Tensor
    def test_check_finite_and_unscale_for_dist_tensor(self):
        x = np.random.random((1024, 1024)).astype("float32")
        x[128][128] = np.inf
        scale = np.random.random(1).astype("float32")
        found_inf = np.array([0]).astype(np.bool_)

        local_x, dist_x = self.create_local_and_dist_tensor_pair(x)
        local_scale, dist_scale = self.create_local_and_dist_tensor_pair(scale)
        (
            local_found_inf,
            dist_found_inf,
        ) = self.create_local_and_dist_tensor_pair(found_inf)

        paddle._C_ops.check_finite_and_unscale_(
            [local_x],
            local_scale,
            [local_x],
            local_found_inf,
        )
        paddle._C_ops.check_finite_and_unscale_(
            [dist_x],
            dist_scale,
            [dist_x],
            dist_found_inf,
        )
        self.check_tensor_eq(local_x, dist_x)
        self.check_tensor_eq(local_found_inf, dist_found_inf)

    # multi kernel functions
    def test_adagrad_for_dist_tensor(self):
        dtype = np.float16
        mp_dtype = np.float32
        shape = [123, 321]

        param = np.random.random(shape).astype(dtype)
        grad = np.random.random(shape).astype(dtype)
        moment = np.random.random(shape).astype(mp_dtype)
        master_param = param.astype(mp_dtype)

        lr = np.array([0.002]).astype("float32")
        epsilon = 1e-8

        local_param, dist_param = self.create_local_and_dist_tensor_pair(param)
        local_grad, dist_grad = self.create_local_and_dist_tensor_pair(grad)
        local_lr, dist_lr = self.create_local_and_dist_tensor_pair(lr)
        local_moment, dist_moment = self.create_local_and_dist_tensor_pair(
            moment
        )
        (
            local_master_param,
            dist_master_param,
        ) = self.create_local_and_dist_tensor_pair(master_param)
        (
            local_param_out,
            local_moment_out,
            local_master_param_out,
        ) = paddle._C_ops.adagrad_(
            local_param,
            local_grad,
            local_moment,
            local_lr,
            local_master_param,
            epsilon,
            True,
        )

        (
            dist_param_out,
            dist_moment_out,
            dist_master_param_out,
        ) = paddle._C_ops.adagrad_(
            dist_param,
            dist_grad,
            dist_moment,
            dist_lr,
            dist_master_param,
            epsilon,
            True,
        )

        self.check_tensor_eq(local_param_out, dist_param_out)
        self.check_tensor_eq(local_moment_out, dist_moment_out)
        self.check_tensor_eq(local_master_param_out, dist_master_param_out)

    # input: std::vector<phi::Tensor>, phi::Tensor
    # output: inplace paddle::optional<std::vector<phi::Tensor>>, inplace phi::Tensor
    def test_merged_adam_for_dist_tensor(self):
        dtype = np.float16
        mp_dtype = np.float32
        lr_shape = [[1], [1], [1], [1]]
        shapes = [[3, 4], [2, 7], [5, 6], [7, 8]]

        epsilon = 0.9
        beta1 = 0.9
        beta2 = 0.99
        params = [np.random.random(s).astype(dtype) for s in shapes]
        grads = [np.random.random(s).astype(dtype) for s in shapes]
        lrs = [np.random.random(s).astype(np.float64) for s in lr_shape]
        moment1s = [np.random.random(s).astype(mp_dtype) for s in shapes]
        moment2s = [np.random.random(s).astype(mp_dtype) for s in shapes]
        moment2s_max = [np.zeros(s).astype(mp_dtype) for s in shapes]
        beta1_pows = [np.random.random(s).astype(mp_dtype) for s in lr_shape]
        beta2_pows = [np.random.random(s).astype(mp_dtype) for s in lr_shape]
        master_params = [p.astype(mp_dtype) for p in params]

        local_param, dist_param = self.create_local_and_dist_tensor_list_pair(
            params
        )
        local_grads, dist_grads = self.create_local_and_dist_tensor_list_pair(
            grads
        )
        local_lrs, dist_lrs = self.create_local_and_dist_tensor_list_pair(lrs)
        (
            local_moment1s,
            dist_moment1s,
        ) = self.create_local_and_dist_tensor_list_pair(moment1s)
        (
            local_moment2s,
            dist_moment2s,
        ) = self.create_local_and_dist_tensor_list_pair(moment2s)
        (
            local_moment2s_max,
            dist_moment2s_max,
        ) = self.create_local_and_dist_tensor_list_pair(moment2s_max)
        (
            local_beta1_pows,
            dist_beta1_pows,
        ) = self.create_local_and_dist_tensor_list_pair(beta1_pows)
        (
            local_beta2_pows,
            dist_beta2_pows,
        ) = self.create_local_and_dist_tensor_list_pair(beta2_pows)
        (
            local_master_params,
            dist_master_params,
        ) = self.create_local_and_dist_tensor_list_pair(master_params)

        (
            local_param_out,
            local_moment1s_out,
            local_moment2s_out,
            local_moment2s_max_out,
            local_beta1_pow_out,
            local_beta2_pow_out,
            local_master_param_out,
        ) = paddle._C_ops.merged_adam_(
            local_param,
            local_grads,
            local_lrs,
            local_moment1s,
            local_moment2s,
            local_moment2s_max,
            local_beta1_pows,
            local_beta2_pows,
            local_master_params,
            beta1,
            beta2,
            epsilon,
            True,
            False,
            False,
        )

        (
            dist_param_out,
            dist_moment1s_out,
            dist_moment2s_out,
            dist_moment2s_max_out,
            dist_beta1_pow_out,
            dist_beta2_pow_out,
            dist_master_param_out,
        ) = paddle._C_ops.merged_adam_(
            dist_param,
            dist_grads,
            dist_lrs,
            dist_moment1s,
            dist_moment2s,
            dist_moment2s_max,
            dist_beta1_pows,
            dist_beta2_pows,
            dist_master_params,
            beta1,
            beta2,
            epsilon,
            True,
            False,
            False,
        )
        for i in range(len(local_param_out)):
            self.check_tensor_eq(local_param_out[i], dist_param_out[i])
            self.check_tensor_eq(local_moment1s_out[i], dist_moment1s_out[i])
            self.check_tensor_eq(local_moment2s_out[i], dist_moment2s_out[i])
            self.check_tensor_eq(local_beta1_pow_out[i], dist_beta1_pow_out[i])
            self.check_tensor_eq(local_beta2_pow_out[i], dist_beta2_pow_out[i])
            self.check_tensor_eq(
                local_master_param_out[i], dist_master_param_out[i]
            )

    # intermediate dygraph api test
    def test_layer_norm_for_intermediate_dist_tensor(self):
        x = np.random.random((2, 3, 10, 10)).astype("float32")
        weight = np.random.random(300).astype("float32")
        bias = np.random.random(300).astype("float32")

        local_x, dist_x = self.create_local_and_dist_tensor_pair(x)
        local_weight, dist_weight = self.create_local_and_dist_tensor_pair(
            weight
        )
        local_bias, dist_bias = self.create_local_and_dist_tensor_pair(bias)

        local_out = paddle.nn.functional.layer_norm(
            local_x,
            local_x.shape[1:],
            local_weight,
            local_bias,
        )
        dist_out = paddle.nn.functional.layer_norm(
            dist_x,
            dist_x.shape[1:],
            dist_weight,
            dist_bias,
        )
        self.check_tensor_eq(local_out, dist_out)


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