# 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.

# coding: utf-8
""" Key value store interface of MXNet for parameter synchronization."""

import pickle
import ctypes
import os
from ..ndarray import NDArray
from ..ndarray import _ndarray_cls
from ..base import _LIB, c_str
from ..base import check_call, mx_uint, py_str
from ..base import NDArrayHandle, KVStoreHandle
from .. import optimizer as opt
from .base import _ctype_key_value, _ctype_dict, KVStoreBase

__all__ = ['KVStore']

def _updater_wrapper(updater):
    """A wrapper for the user-defined handle."""
    def updater_handle(key, lhs_handle, rhs_handle, _):
        """ ctypes function """
        lhs = _ndarray_cls(NDArrayHandle(lhs_handle))
        rhs = _ndarray_cls(NDArrayHandle(rhs_handle))
        updater(key, lhs, rhs)
    return updater_handle

def _get_kvstore_server_command_type(command):
    command_types = {'kController': 0,
                     'kSetMultiPrecision': 1,
                     'kStopServer': 2,
                     'kSyncMode': 3,
                     'kSetGradientCompression': 4,
                     'kSetProfilerParams': 5}
    assert (command in command_types), "Unknown command type to send to server"
    return command_types[command]


class KVStore(KVStoreBase):
    """A key-value store for synchronization of values, over multiple devices."""

    def __init__(self, handle):
        """Initializes a new KVStore.

        Parameters
        ----------
        handle : KVStoreHandle
            `KVStore` handle of C API.
        """
        assert isinstance(handle, KVStoreHandle)
        self.handle = handle
        self._updater = None
        self._updater_func = None
        self._str_updater_func = None
        self._is_p3 = (os.getenv('DMLC_PS_VAN_TYPE', '') == 'p3')

    def __del__(self):
        check_call(_LIB.MXKVStoreFree(self.handle))

    def broadcast(self, key, value, out, priority=0):
        """ Broadcast the `value` NDArray at rank 0 to all ranks,
        and store the result in `out`.

        Note that the native KVStore does not support broadcasting the same key more than once.

        Parameters
        ----------
        key : str, or int
            The key.

        value : NDArray, list of NDArray, or list of list of NDArray
            Values corresponding to the keys.

        out: NDArray or list of NDArray or list of list of NDArray
            Outputs corresponding to the keys.

        priority : int, optional
            The priority of the operation.
            Higher priority operations are likely to be executed before other actions.

        Examples
        --------
        >>> # broadcast a single key-value pair
        >>> shape = (2,3)
        >>> kv = mx.kv.create('local')
        >>> a = mx.nd.zeros(shape)
        >>> kv.broadcast('3', mx.nd.ones(shape)*2, out=a)
        >>> print a.asnumpy()
        [[ 2.  2.  2.]
        [ 2.  2.  2.]]

        """
        cvkeys, cvals, use_str_keys = _ctype_key_value(key, value)
        cokeys, couts, _ = _ctype_key_value(key, out)

        if use_str_keys:
            check_call(_LIB.MXKVStoreBroadcastEx(
                self.handle, mx_uint(len(cvkeys)), cvkeys, mx_uint(len(cokeys)), cokeys,
                cvals, couts, ctypes.c_int(priority)))
        else:
            check_call(_LIB.MXKVStoreBroadcast(
                self.handle, mx_uint(len(cvkeys)), cvkeys, mx_uint(len(cokeys)), cokeys,
                cvals, couts, ctypes.c_int(priority)))


    def is_capable(self, capability):
        """Queries if the KVStore type supports certain capability, such as optimizer algorithm,
        gradient compression, sparsity, etc.

        Parameters
        ----------
        capability: str
            The capability to query

        Returns
        -------
        result : bool
            Whether the capability is supported or not.
        """
        if capability.lower() == KVStoreBase.OPTIMIZER:
            return not self._is_p3
        else:
            raise ValueError('Unknown capability: {}'.format(capability))

    def init(self, key, value):
        """ Initializes a single or a sequence of key-value pairs into the store.

        For each key, one must `init` it before calling `push` or `pull`.
        When multiple workers invoke `init` for the same key, only
        the value supplied by worker with rank `0` is used. This function returns
        after data has been initialized successfully.

        Parameters
        ----------
        key : str, int, or sequence of str or int
            The keys.
        value : NDArray, RowSparseNDArray or sequence of NDArray or RowSparseNDArray
            Values corresponding to the keys.

        Examples
        --------
        >>> # init a single key-value pair
        >>> shape = (2,3)
        >>> kv = mx.kv.create('local')
        >>> kv.init('3', mx.nd.ones(shape)*2)
        >>> a = mx.nd.zeros(shape)
        >>> kv.pull('3', out=a)
        >>> print a.asnumpy()
        [[ 2.  2.  2.]
        [ 2.  2.  2.]]

        >>> # init a list of key-value pairs
        >>> keys = ['5', '7', '9']
        >>> kv.init(keys, [mx.nd.ones(shape)]*len(keys))

        >>> # init a row_sparse value
        >>> kv.init('4', mx.nd.ones(shape).tostype('row_sparse'))
        >>> b = mx.nd.sparse.zeros('row_sparse', shape)
        >>> kv.row_sparse_pull('4', row_ids=mx.nd.array([0, 1]), out=b)
        >>> print b
        <RowSparseNDArray 2x3 @cpu(0)>
        """
        ckeys, cvals, use_str_keys = _ctype_key_value(key, value)
        if use_str_keys:
            check_call(_LIB.MXKVStoreInitEx(self.handle, mx_uint(len(ckeys)), ckeys, cvals))
        else:
            check_call(_LIB.MXKVStoreInit(self.handle, mx_uint(len(ckeys)), ckeys, cvals))

    def push(self, key, value, priority=0):
        """ Pushes a single or a sequence of key-value pairs into the store.

        This function returns immediately after adding an operator to the engine.
        The actual operation is executed asynchronously. If there are consecutive
        pushes to the same key, there is no guarantee on the serialization of pushes.
        The execution of a push does not guarantee that all previous pushes are
        finished.
        There is no synchronization between workers.
        One can use ``_barrier()`` to sync all workers.

        Parameters
        ----------
        key : str, int, or sequence of str or int
            Keys.

        value : NDArray, RowSparseNDArray, list of NDArray or RowSparseNDArray,
                or list of list of NDArray or RowSparseNDArray
            Values corresponding to the keys.

        priority : int, optional
            The priority of the push operation.
            Higher priority push operations are likely to be executed before
            other push actions.

        Examples
        --------
        >>> # push a single key-value pair
        >>> shape = (2,3)
        >>> kv.push('3', mx.nd.ones(shape)*8)
        >>> kv.pull('3', out=a) # pull out the value
        >>> print a.asnumpy()
        [[ 8.  8.  8.]
        [ 8.  8.  8.]]

        >>> # aggregate the value and the push
        >>> gpus = [mx.gpu(i) for i in range(4)]
        >>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]
        >>> kv.push('3', b)
        >>> kv.pull('3', out=a)
        >>> print a.asnumpy()
        [[ 4.  4.  4.]
        [ 4.  4.  4.]]

        >>> # push a list of keys.
        >>> # single device
        >>> keys = ['4', '5', '6']
        >>> kv.push(keys, [mx.nd.ones(shape)]*len(keys))
        >>> b = [mx.nd.zeros(shape)]*len(keys)
        >>> kv.pull(keys, out=b)
        >>> print b[1].asnumpy()
        [[ 1.  1.  1.]
        [ 1.  1.  1.]]

        >>> # multiple devices:
        >>> keys = ['7', '8', '9']
        >>> b = [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys)
        >>> kv.push(keys, b)
        >>> kv.pull(keys, out=b)
        >>> print b[1][1].asnumpy()
        [[ 4.  4.  4.]
        [ 4.  4.  4.]]

        >>> # push a row_sparse value
        >>> b = mx.nd.sparse.zeros('row_sparse', shape)
        >>> kv.init('10', mx.nd.sparse.zeros('row_sparse', shape))
        >>> kv.push('10', mx.nd.ones(shape).tostype('row_sparse'))
        >>> # pull out the value
        >>> kv.row_sparse_pull('10', row_ids=mx.nd.array([0, 1]), out=b)
        >>> print b
        <RowSparseNDArray 2x3 @cpu(0)>
        """
        ckeys, cvals, use_str_keys = _ctype_key_value(key, value)
        if use_str_keys:
            check_call(_LIB.MXKVStorePushEx(
                self.handle, mx_uint(len(ckeys)), ckeys, cvals, ctypes.c_int(priority)))
        else:
            check_call(_LIB.MXKVStorePush(
                self.handle, mx_uint(len(ckeys)), ckeys, cvals, ctypes.c_int(priority)))


    def pull(self, key, out=None, priority=0, ignore_sparse=True):
        """ Pulls a single value or a sequence of values from the store.

        This function returns immediately after adding an operator to the engine.
        Subsequent attempts to read from the `out` variable will be blocked until the
        pull operation completes.

        `pull` is executed asynchronously after all previous `pull` calls and only
        the last `push` call for the same input key(s) are finished.

        The returned values are guaranteed to be the latest values in the store.

        pull with `RowSparseNDArray` is not supported for dist kvstore.
        Please use ``row_sparse_pull`` instead.

        Parameters
        ----------
        key : str, int, or sequence of str or int
            Keys.

        out: NDArray or list of NDArray or list of list of NDArray