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

package AI::MXNet::KVStore;
use strict;
use warnings;
use AI::MXNet::Base;
use AI::MXNet::NDArray;
use AI::MXNet::Optimizer;
use MIME::Base64;
use Storable;
use Mouse;
use AI::MXNet::Function::Parameters;

=head1 NAME

    AI::MXNet::KVStore - Key value store interface of MXNet.

=head1 DESCRIPTION

    Key value store interface of MXNet for parameter synchronization, over multiple devices.
=cut

has 'handle' => (is => 'ro', isa => 'KVStoreHandle', required => 1);
has '_updater' => (is => 'rw',  isa => 'AI::MXNet::Updater');

sub DEMOLISH
{
    check_call(AI::MXNetCAPI::KVStoreFree(shift->handle));
}

=head2  init

    Initialize a single or a sequence of key-value pairs into the store.
    For each key, one must init it before push and pull.
    Only worker 0's (rank == 0) data are used.
    This function returns after data have been initialized successfully

    Parameters
    ----------
    $key : Str|ArrayRef[Str]
        The keys.
    $value : AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]]
        The values.

    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->aspdl
    [[ 2  2  2]
    [ 2  2  2]]

    >>> # init a list of key-value pairs
    >>> $keys = [5, 7, 9]
    >>> $kv->init(keys, [map { mx->nd->ones($shape) } 0..@$keys-1])
=cut

method init(
    Str|ArrayRef[Str] $key,
    AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] $value
)
{
    my ($keys, $vals) = _key_value($key, $value);
    check_call(
        AI::MXNetCAPI::KVStoreInitEx(
            $self->handle, scalar(@{ $keys }), $keys, $vals
        )
    );
}

=head2  push

    Push a single or a sequence of key-value pairs into the store.
    Data consistency:
    1. this function returns after adding an operator to the engine.
    2. push is always called after all previous push and pull on the same
        key are finished.
    3. there is no synchronization between workers. One can use _barrier()
    to sync all workers.

    Parameters
    ----------
    $key : Str|ArrayRef[Str]
    $value : AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]]
    :$priority=0 : Int, optional
        The priority of the push operation.
        The higher the priority, the faster this action is likely
        to be executed before other push actions.

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

    >>> # aggregate the value and the push
    >>> $gpus = [map { mx->gpu($_) } 0..3]
    >>> $b = [map { mx->nd->ones($shape, ctx => $_) } @$gpus]
    >>> $kv->push(3, $b)
    >>> $kv->pull(3, out=>$a)
    >>> print $a->aspdl
        [[ 4.  4.  4.]
        [ 4.  4.  4.]]

    >>> # push a list of keys.
    >>> # single device
    >>> $kv->push($keys, [map { mx->nd->ones($shape) } 0..@$keys-1)
    >>> $b = [map { mx->nd->zeros(shape) } 0..@$keys-1]
    >>> $kv->pull($keys, out=>$b)
    >>> print $b->[1]->aspdl
        [[ 1.  1.  1.]
        [ 1.  1.  1.]]

    >>> # multiple devices:
    >>> $b = [map { [map { mx->nd->ones($shape, ctx => $_) } @$gpus] } @$keys-1]
    >>> $kv->push($keys, $b)
    >>> $kv->pull($keys, out=>$b)
    >>> print $b->[1][1]->aspdl()
        [[ 4.  4.  4.]
        [ 4.  4.  4.]]
=cut

method push(
    Str|ArrayRef[Str] $key,
    AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] $value,
    Int :$priority=0
)
{
    my ($keys, $vals) = _key_value($key, $value);
    check_call(
        AI::MXNetCAPI::KVStorePushEx(
            $self->handle, scalar(@{ $keys }), $keys, $vals, $priority
        )
    );
}

=head2 pull

    Pull a single value or a sequence of values from the store.

    Data consistency:

    1. this function returns after adding an operator to the engine. But any
        further read on out will be blocked until it is finished.
    2. pull is always called after all previous push and pull on the same
        key are finished.
    3. It pulls the newest value from the store.

    Parameters
    ----------
    $key : Str|ArrayRef[Str]
        Keys
    :$out: AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]]
        According values

    :$priority=0 : Int, optional
        The priority of the push operation.
        The higher the priority, the faster this action is likely
        to be executed before other push actions.

    Examples
    --------
    >>> # pull a single key-value pair
    >>> $a = mx->nd->zeros($shape)
    >>> $kv->pull(3, out=>$a)
    >>> print $a->aspdl
        [[ 2.  2.  2.]
        [ 2.  2.  2.]]

    >>> # pull into multiple devices
    >>> $b = [map { mx->nd->ones($shape, $_) } @$gpus]
    >>> $kv->pull(3, out=>$b)
    >>> print $b->[1]->aspdl()
        [[ 2.  2.  2.]
        [ 2.  2.  2.]]

    >>> # pull a list of key-value pairs.
    >>> # On single device
    >>> $keys = [5, 7, 9]
    >>> $b = [map { mx->nd->zeros($shape) } 0..@$keys-1]
    >>> $kv->pull($keys, out=>$b)
    >>> print $b->[1]->aspdl()
        [[ 2.  2.  2.]
        [ 2.  2.  2.]]
    >>> # On multiple devices
    >>> $b = [map { [map { mx->nd->ones($shape, ctx => $_) } @$gpus ] } 0..@$keys-1]
    >>> $kv->pull($keys, out=>$b)
    >>> print $b->[1][1]->aspdl()
        [[ 2.  2.  2.]
        [ 2.  2.  2.]]
=cut

method pull(
    Str|ArrayRef[Str] $key,
    AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] :$out,
    Int :$priority=0
)
{
    my ($keys, $vals) = _key_value($key, $out);
    check_call(
        AI::MXNetCAPI::KVStorePullEx(
            $self->handle, scalar(@{ $keys }), $keys, $vals, $priority
        )
    );
}

=head2  row_sparse_pull

        Pulls a single AI::MXNet::NDArray::RowSparse value or an array ref of AI::MXNet::NDArray::RowSparse values
        from the store with specified row_ids. When there is only one row_id, KVStoreRowSparsePull
        is invoked just once and the result is broadcast to all the rest of outputs.

        `row_sparse_pull` is executed asynchronously after all previous
        `pull`/`row_sparse_pull` calls and 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.

        Parameters
        ----------
        $key : Str|ArrayRef[Str] $key
            Keys.

        :$out: AI::MXNet::NDArray::RowSparse|ArrayRef[AI::MXNet::NDArray::RowSparse]|ArrayRef[ArrayRef[AI::MXNet::NDArray::RowSparse]]
            Values corresponding to the keys. The stype is expected to be row_sparse

        :$priority=0 : Int, optional
            The priority of the pull operation.
            Higher priority pull operations are likely to be executed before
            other pull actions.

        :$row_ids : AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]]
            The row_ids for which to pull for each value. Each row_id is an 1D NDArray
            whose values don't have to be unique nor sorted.

        Examples
        --------
        >>> $shape = [3, 3]
        >>> $kv->init('3', mx->nd->ones($shape)->tostype('row_sparse'))
        >>> $a = mx->nd->sparse->zeros('row_sparse', $shape)
        >>> $row_ids = mx->nd->array([0, 2], dtype=>'int64')
        >>> $kv->row_sparse_pull('3', out=>$a, row_ids=>$row_ids)
        >>> print $a->aspdl
        [[ 1.  1.  1.]
        [ 0.  0.  0.]
        [ 1.  1.  1.]]
        >>> $duplicate_row_ids = mx->nd->array([2, 2], dtype=>'int64')
        >>> $kv->row_sparse_pull('3', out=>$a, row_ids=>$duplicate_row_ids)
        >>> print $a->aspdl
        [[ 0.  0.  0.]
        [ 0.  0.  0.]
        [ 1.  1.  1.]]
        >>> $unsorted_row_ids = mx->nd->array([1, 0], dtype=>'int64')
        >>> $kv->row_sparse_pull('3', out=>$a, row_ids=>$unsorted_row_ids)
        >>> print $a->aspdl
        [[ 1.  1.  1.]
        [ 1.  1.  1.]
        [ 0.  0.  0.]]
=cut


method row_sparse_pull(
    Str|ArrayRef[Str] $key,
    AI::MXNet::NDArray::RowSparse|ArrayRef[AI::MXNet::NDArray::RowSparse]|ArrayRef[ArrayRef[AI::MXNet::NDArray::RowSparse]] :$out,
    Int :$priority=0,
    AI::MXNet::NDArray|ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] :$row_ids
)
{
    if(blessed $row_ids)
    {
        $row_ids = [$row_ids];
    }
    my $first_out = $out;
    # whether row_ids are the same
    my $single_rowid = 0;
    if(@$row_ids == 1 and ref $out eq 'ARRAY')
    {
        $single_rowid = 1;
        $first_out = [$out->[0]];
    }
    my ($ckeys, $cvals) = _key_value($key, $first_out);
    my (undef, $crow_ids) = _key_value($key, $row_ids);
    assert(
        (@$crow_ids == @$cvals),
        "the number of row_ids doesn't match the number of values"
    );
    check_call(
        AI::MXNetCAPI::KVStorePullRowSparseEx(
            $self->handle, scalar(@$ckeys), $ckeys, $cvals, $crow_ids, $priority
        )
    );
    # the result can be copied to other devices without invoking row_sparse_pull
    # if the indices are the same
    if($single_rowid)
    {
        for my $out_i (@{ $out } [1..@{ $out }-1])
        {
            $out->[0]->copyto($out_i);
        }
    }
}

=head2  set_gradient_compression

        Specifies type of low-bit quantization for gradient compression \
         and additional arguments depending on the type of compression being used.

        2bit Gradient Compression takes a positive float `threshold`.
        The technique works by thresholding values such that positive values in the
        gradient above threshold will be set to threshold. Negative values whose absolute
        values are higher than threshold, will be set to the negative of threshold.
        Values whose absolute values are less than threshold will be set to 0.
        By doing so, each value in the gradient is in one of three states. 2bits are
        used to represent these states, and every 16 float values in the original
        gradient can be represented using one float. This compressed representation
        can reduce communication costs. The difference between these thresholded values and
        original values is stored at the sender's end as residual and added to the
        gradient in the next iteration.

        When kvstore is 'local', gradient compression is used to reduce communication
        between multiple devices (gpus). Gradient is quantized on each GPU which
        computed the gradients, then sent to the GPU which merges the gradients. This
        receiving GPU dequantizes the gradients and merges them. Note that this
        increases memory usage on each GPU because of the residual array stored.

        When kvstore is 'dist', gradient compression is used to reduce communication
        from worker to sender. Gradient is quantized on each worker which
        computed the gradients, then sent to the server which dequantizes
        this data and merges the gradients from each worker. Note that this
        increases CPU memory usage on each worker because of the residual array stored.
        Only worker to server communication is compressed in this setting.
        If each machine has multiple GPUs, currently this GPU to GPU or GPU to CPU communication
        is not compressed. Server to worker communication (in the case of pull)
        is also not compressed.

        To use 2bit compression, we need to specify `type` as `2bit`.
        Only specifying `type` would use default value for the threshold.
        To completely specify the arguments for 2bit compression, we would need to pass
        a dictionary which includes `threshold` like:
        {'type': '2bit', 'threshold': 0.5}

        Parameters
        ----------
        $compression_params : HashRef[Str]
            A dictionary specifying the type and parameters for gradient compression.
            The key `type` in this dictionary is a
            required string argument and specifies the type of gradient compression.
            Currently `type` can be only `2bit`
            Other keys in this dictionary are optional and specific to the type
            of gradient compression.
=cut

method set_gradient_compression(HashRef[Str] $compression_params)
{
    if($self->type =~ /(?:device|dist)/)
    {
        check_call(
            AI::MXNetCAPI::KVStoreSetGradientCompression(
                $self->handle,
                scalar(keys %$compression_params),
                $compression_params
            )
        );
    }
    else
    {
        confess('Gradient compression is not supported for this type of kvstore');
    }
}

=head2  set_optimizer

    Register an optimizer to the store

    If there are multiple machines, this process (should be a worker node)
    will pack this optimizer and send it to all servers. It returns after
    this action is done.

    Parameters
    ----------
    $optimizer : AI::MXNet::Optimizer
        the optimizer
=cut

method set_optimizer(AI::MXNet::Optimizer $optimizer)
{
    my $is_worker = check_call(AI::MXNetCAPI::KVStoreIsWorkerNode());
    if($self->type =~ /dist/ and $is_worker)
    {
        my $optim_str = MIME::Base64::encode_base64(Storable::freeze($optimizer), "");
        $self->_send_command_to_servers(0, $optim_str);
    }
    else
    {
        $self->_updater(AI::MXNet::Optimizer->get_updater($optimizer));
        $self->_set_updater($self->_updater);
    }
}

=head2  type

    Get the type of this kvstore

    Returns
    -------
    $type : Str
        the string type
=cut

method type()
{
    return scalar(check_call(AI::MXNetCAPI::KVStoreGetType($self->handle)));
}

=head2  rank

    Get the rank of this worker node

    Returns
    -------
    $rank : Int
        The rank of this node, which is in [0, get_num_workers())
=cut

method rank()
{
    return scalar(check_call(AI::MXNetCAPI::KVStoreGetRank($self->handle)));
}

=head2  num_workers

    Get the number of worker nodes

    Returns
    -------
    $size : Int
        The number of worker nodes
=cut

method num_workers()
{
    return scalar(check_call(AI::MXNetCAPI::KVStoreGetGroupSize($self->handle)));
}

=head2 save_optimizer_states

    Save optimizer (updater) state to file

    Parameters
    ----------
    $fname : Str
        Path to output states file.
    :$dump_optimizer=0 : Bool, default False
            Whether to also save the optimizer itself. This would also save optimizer
            information such as learning rate and weight decay schedules.
=cut

method save_optimizer_states(Str $fname, Bool :$dump_optimizer=0)
{
    confess("Cannot save states for distributed training")
        unless defined $self->_updater;
    open(F, ">:raw", "$fname") or confess("can't open $fname for writing: $!");
    print F $self->_updater->get_states($dump_optimizer);
    close(F);
}

=head2 load_optimizer_states

    Load optimizer (updater) state from file.

    Parameters
    ----------
    $fname : Str
        Path to input states file.
=cut

method load_optimizer_states(Str $fname)
{
    confess("Cannot save states for distributed training")
        unless defined $self->_updater;
    open(F, "<:raw", "$fname") or confess("can't open $fname for reading: $!");
    my $data;
    { local($/) = undef; $data = <F>; }
    close(F);
    $self->_updater->set_states($data);
}

=head2 _set_updater

    Set a push updater into the store.

    This function only changes the local store. Use set_optimizer for
    multi-machines.

    Parameters
    ----------
    $updater : Undater
        the updater function

    Examples
    --------
    >>> my $update = sub { my ($key, input, stored) = @_;
        ...     print "update on key: $key\n";
        ...     $stored += $input * 2; };
        >>> $kv->_set_updater($update)
        >>> $kv->pull(3, out=>$a)
        >>> print $a->aspdl()
        [[ 4.  4.  4.]
        [ 4.  4.  4.]]
        >>> $kv->push(3, mx->nd->ones($shape))
        update on key: 3
        >>> $kv->pull(3, out=>$a)
        >>> print $a->aspdl()
        [[ 6.  6.  6.]
        [ 6.  6.  6.]]
=cut

method _set_updater(Updater $updater_func)
{
    check_call(
        AI::MXNetCAPI::KVStoreSetUpdater(
            $self->handle,
            sub {
                my ($index, $input_handle, $storage_handle) = @_;
                $updater_func->(
                    $index,
                    AI::MXNet::NDArray->_ndarray_cls($input_handle),
                    AI::MXNet::NDArray->_ndarray_cls($storage_handle)
                );
            }
        )
    );
}

=head2 _barrier

    Global barrier between all worker nodes.

    For example, assume there are n machines, we want to let machine 0 first
    init the values, and then pull the inited value to all machines. Before
    pulling, we can place a barrier to guarantee that the initialization is
    finished.
=cut

method _barrier()
{
    check_call(AI::MXNetCAPI::KVStoreBarrier($self->handle));
}

=head2 _send_command_to_servers

    Send a command to all server nodes
    Send a command to all server nodes, which will make each server node run
    KVStoreServer.controller
    This function returns after the command has been executed in all server
    nodes.

    Parameters
    ----------
    $head : Int
        the head of the command
    $body : Str
        the body of the command
=cut

method _send_command_to_servers(Int $head, Str $body)
{
    check_call(
        AI::MXNetCAPI::KVStoreSendCommmandToServers(
            $self->handle,
            $head,
            $body
        )
    );
}

=head2 create

    Create a new KVStore.

    Parameters
    ----------
    $name='local' : Str
    The type of KVStore
        - local works for multiple devices on a single machine (single process)
        - dist works for multi-machines (multiple processes)
    Returns
    -------
    kv : KVStore
        The created AI::MXNet::KVStore
=cut

method create(Str $name='local')
{
    my $handle = check_call(AI::MXNetCAPI::KVStoreCreate($name));
    return __PACKAGE__->new(handle => $handle);
}

sub _key_value
{
    my ($keys, $vals) = @_;
    if(not ref $keys)
    {
        if(blessed $vals)
        {
            return ([$keys], [$vals->handle]);
        }
        else
        {
            for my $value (@{ $vals })
            {
                assert(blessed($value) and $value->isa('AI::MXNet::NDArray'));
                return ([($keys)x@$vals], [map { $_->handle } @$vals]);
            }
        }
    }
    else
    {
        assert(not blessed($vals) and @$keys == @$vals);
        my @c_keys;
        my @c_vals;
        for(zip($keys, $vals)) {
            my ($key, $val) = @$_;
            my ($c_key, $c_val) = _key_value($key, $val);
            push @c_keys, @$c_key;
            push @c_vals, @$c_val;
        }
        return (\@c_keys, \@c_vals);
    }
}

1;
