# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import numpy as np

from ..functional import param_pack_concat, param_pack_split
from ..functional.distributed import all_reduce_sum
from ..tensor import Tensor


def get_offsets(shapes):
    offsets = []
    offset = 0
    for shape in shapes:
        offsets.append(offset)
        offset += int(np.prod(shape))
        offsets.append(offset)
    return offsets


def get_pack_list(param_group, param_pack_thd):
    pack_list = dict()
    shape_list = dict()
    pack_sum = dict()
    pack_ret, shape_ret = [], []
    ignore_first = 8
    ignore_last = 0
    orders_len = len(param_group["orders"])
    for i, idx in enumerate(param_group["orders"]):
        param = param_group["params"][idx]
        dtype = str(np.dtype(param.dtype))
        dtype_size = np.dtype(param.dtype).itemsize
        shape = param.shape
        if ignore_first > 0:
            ignore_first -= 1
            pack_ret.append([idx])
            shape_ret.append([shape])
            continue
        if dtype in pack_list.keys():
            pack_list[dtype].append(idx)
            shape_list[dtype].append(shape)
            pack_sum[dtype] += int(np.prod(shape))
        else:
            pack_list[dtype] = [idx]
            shape_list[dtype] = [shape]
            pack_sum[dtype] = int(np.prod(shape))
        if (
            pack_sum[dtype] * dtype_size > param_pack_thd
            or i + ignore_last > orders_len
        ):
            pack_ret.append(pack_list[dtype])
            shape_ret.append(shape_list[dtype])
            pack_list[dtype] = []
            shape_list[dtype] = []
            pack_sum[dtype] = 0
    for key in sorted(pack_list.keys()):
        if len(pack_list[key]) > 0:
            pack_ret.append(pack_list[key])
            shape_ret.append(shape_list[key])
    return pack_ret, shape_ret


def pack_allreduce_split(group, pack, shapes, reduce_method):
    dist_group = group["dist_group"]
    grads = [group["grads"][idx] for idx in pack]
    offsets_val = get_offsets(shapes)
    offsets = Tensor(offsets_val)
    packed_grads = param_pack_concat(grads, offsets, offsets_val)
    packed_grads = all_reduce_sum(packed_grads, dist_group, dist_group.comp_node)
    if reduce_method == "mean":
        packed_grads /= dist_group.size
    grads = param_pack_split(packed_grads, offsets_val, shapes)
    for i, grad in enumerate(grads):
        group["grads"][pack[i]] = grad