# Copyright (c) 2021 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 random
import unittest

import numpy as np

import paddle
import paddle.distributed as dist
from paddle.distributed import fleet


def set_random_seed(seed):
    """Set random seed for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    paddle.seed(seed)
    fleet.meta_parallel.model_parallel_random_seed(seed)


def class_center_sample_numpy(label, classes_list, num_samples):
    unique_label = np.unique(label)
    nranks = len(classes_list)
    class_interval = np.cumsum(np.insert(classes_list, 0, 0))
    pos_class_center_per_device = []
    unique_label_per_device = []

    for i in range(nranks):
        index = np.logical_and(
            unique_label >= class_interval[i],
            unique_label < class_interval[i + 1],
        )
        pos_class_center_per_device.append(
            unique_label[index] - class_interval[i]
        )
        unique_label_per_device.append(unique_label[index])

    num_samples_per_device = []
    for pos_class_center in pos_class_center_per_device:
        num_samples_per_device.append(max(len(pos_class_center), num_samples))
    sampled_class_interval = np.cumsum(np.insert(num_samples_per_device, 0, 0))

    remapped_dict = {}
    for i in range(nranks):
        for idx, v in enumerate(
            unique_label_per_device[i], sampled_class_interval[i]
        ):
            remapped_dict[v] = idx

    remapped_label = []
    for l in label:
        remapped_label.append(remapped_dict[l])

    return remapped_label, pos_class_center_per_device


class TestParallelClassCenterSampleOp(unittest.TestCase):
    def setUp(self):
        strategy = fleet.DistributedStrategy()
        fleet.init(is_collective=True, strategy=strategy)

    def test_class_center_sample(self):
        rank_id = dist.get_rank()
        nranks = dist.get_world_size()

        seed = 1025
        set_random_seed(seed)
        paddle.seed(rank_id * 10)
        random.seed(seed)
        np.random.seed(seed)

        batch_size = 20
        num_samples = 6

        for dtype in ('int32', 'int64'):
            for _ in range(5):
                classes_list = np.random.randint(10, 15, (nranks,))
                num_class = np.sum(classes_list)

                np_label = np.random.randint(
                    0, num_class, (batch_size,), dtype=dtype
                )
                label = paddle.to_tensor(np_label, dtype=dtype)
                (
                    np_remapped_label,
                    np_sampled_class_center_per_device,
                ) = class_center_sample_numpy(
                    np_label, classes_list, num_samples
                )
                (
                    remapped_label,
                    sampled_class_index,
                ) = paddle.nn.functional.class_center_sample(
                    label, classes_list[rank_id], num_samples
                )
                np.testing.assert_allclose(
                    remapped_label.numpy(), np_remapped_label
                )
                np_sampled_class_index = np_sampled_class_center_per_device[
                    rank_id
                ]
                np.testing.assert_allclose(
                    sampled_class_index.numpy()[: len(np_sampled_class_index)],
                    np_sampled_class_index,
                )


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