# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

# pylint: skip-file
import mxnet as mx
import numpy as np
import unittest
from mxnet.test_utils import rand_ndarray, assert_almost_equal
from common import assertRaises
from mxnet.base import py_str, MXNetError

shape = (4, 4)
keys = [5, 7, 11]
str_keys = ['b', 'c', 'd']

def check_diff_to_scalar(A, x):
    """ assert A == x"""
    assert(np.sum(np.abs((A - x).asnumpy())) == 0), (A, x)

def init_kv(name='device'):
    return mx.kv.create(name)

def test_broadcast_single_kv_pair():
    """single key-value pair push & pull"""
    def check_single_kv_pair(kv, key):
        # single output
        ones = mx.nd.ones(shape)
        out = mx.nd.empty(shape)
        kv.broadcast(key, ones, out)
        check_diff_to_scalar(out, 1)
        # list output
        out_list = [mx.nd.empty(shape)] * 3
        key_list = key + key
        kv.broadcast(key_list, ones, out_list)
        for o in out_list:
            check_diff_to_scalar(o, 1)

    for name in ['device', 'teststore']:
        check_single_kv_pair(init_kv(name), 3)
        check_single_kv_pair(init_kv(name), 'a')

def test_broadcast_list_kv_pair():
    """list key-value pair push & pull"""
    def check_list_kv_pair(kv, key):
        ones = [mx.nd.ones(shape)] * len(key)
        out = [mx.nd.empty(shape)] * len(key)
        kv.broadcast(key, ones, out)
        for o in out:
            check_diff_to_scalar(o, 1)
        out_list = [[mx.nd.empty(shape)] * 2 for _ in range(len(key))]
        key_list = [k + k for k in key]
        kv.broadcast(key_list, ones, out_list)
        for o in out_list:
            for oo in o:
                check_diff_to_scalar(oo, 1)

    check_list_kv_pair(init_kv(), keys)
    check_list_kv_pair(init_kv(), str_keys)

def test_pushpull_single_kv_pair():
    """aggregate value on muliple devices"""
    def check_aggregator(kv, key, key_list=None):
        kv.broadcast(key, mx.nd.zeros(shape), out=mx.nd.empty(shape))
        # devices
        num_devs = 4
        devs = [mx.Context('cpu', i) for i in range(num_devs)]

        # single
        vals = [mx.nd.ones(shape, d) for d in devs]
        outs = [mx.nd.empty(shape, d) for d in devs]

        kv.pushpull(key, vals, out=outs)
        for out in outs:
            check_diff_to_scalar(out, num_devs)

        # inplace
        kv.pushpull(key, vals)
        for val in vals:
            check_diff_to_scalar(val, num_devs)

        # list
        if key_list is None:
            return
        num_keys = len(key_list)
        kv.broadcast(key_list, [mx.nd.zeros(shape)] * num_keys,
                     out=[mx.nd.empty(shape)] * num_keys)
        vals = [[mx.nd.ones(shape, d)*2.0 for d in devs]] * num_keys
        outs = [[mx.nd.empty(shape, d) for d in devs]] * num_keys
        kv.pushpull(key_list, vals, out=outs)
        for out in outs:
            for o in out:
                check_diff_to_scalar(o, num_devs * 2.0)

        # inplace
        kv.pushpull(key_list, vals)
        for val in vals:
            for v in val:
                check_diff_to_scalar(v, num_devs * 2.0)

    check_aggregator(init_kv('device'), 3, keys)
    check_aggregator(init_kv('device'), 'a', str_keys)
    check_aggregator(init_kv('teststore'), 3)
    check_aggregator(init_kv('teststore'), 'a')

def test_pushpull_list_kv_pair():
    """aggregate value on muliple devices"""
    def check_aggregator(kv, key, key_list=None):
        kv.broadcast(key, mx.nd.zeros(shape), out=mx.nd.empty(shape))
        # devices
        num_devs = 4
        devs = [mx.Context('cpu', i) for i in range(num_devs)]

        # single
        vals = [mx.nd.ones(shape, d) for d in devs]
        outs = [mx.nd.empty(shape, d) for d in devs]

        kv.pushpull(key, vals, out=outs)
        for out in outs:
            check_diff_to_scalar(out, num_devs)

        # list
        if key_list is None:
            return
        num_keys = len(key_list)
        kv.broadcast(key_list, [mx.nd.zeros(shape)] * num_keys,
                     out=[mx.nd.empty(shape)] * num_keys)
        vals = [[mx.nd.ones(shape, d)*2.0 for d in devs]] * num_keys
        outs = [[mx.nd.empty(shape, d) for d in devs]] * num_keys
        kv.pushpull(key_list, vals, out=outs)
        for out in outs:
            for o in out:
                check_diff_to_scalar(o, num_devs * 2.0)

    check_aggregator(init_kv('device'), 3, keys)
    check_aggregator(init_kv('device'), 'a', str_keys)
    check_aggregator(init_kv('teststore'), 3)
    check_aggregator(init_kv('teststore'), 'a')


def test_custom_store():
    kv = mx.kv.create('teststore')
    out = mx.nd.empty((1,))
    kv.broadcast(1, mx.nd.ones((1,)), out=out)
    check_diff_to_scalar(out, 1)
    assert type(kv).is_capable('optimizer') == False
    kv.broadcast(1, mx.nd.ones((1,)), out=out)
    check_diff_to_scalar(out, 1)
    arr_list = [mx.nd.empty((1,))] * 2
    kv.pushpull(1, [mx.nd.ones((1,))] * 2, out=arr_list)
    for arr in arr_list:
        check_diff_to_scalar(arr, 2)
    kv.pushpull(1, arr_list)
    for arr in arr_list:
        check_diff_to_scalar(arr, 4)

def test_get_type_device():
    kvtype = 'teststore'
    kv = mx.kv.create(kvtype)
    assert kv.type == kvtype

def test_set_optimizer():
    def check_unsupported_methods(kv):
        assert not kv.is_capable('optimizer')
        optimizer = mx.optimizer.create('sgd')
        assertRaises(NotImplementedError, kv.set_optimizer, optimizer)
        assertRaises(NotImplementedError, kv.save_optimizer_states, 'test')
        assertRaises(NotImplementedError, kv.load_optimizer_states, 'test')

    kv = mx.kv.create('teststore')
    check_unsupported_methods(kv)