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

from mxnet.ndarray import NDArray, topk, abs as NDabs
from mxnet.optimizer import SGD, register
import logging

log = 'Sparsity Update:\t'

@register
class SparseSGD(SGD):
    """The SGD optimizer with weight pruning.

    This class implements the optimizer described in the paper *DSD: Dense-Sparse-Dense Training for
    Deep Neural Networks*, available at https://arxiv.org/pdf/1607.04381.pdf

    The optimizer updates the weights the same way as done in SGD, but does the following
    preprocessing::

        if threshold given, all weights below the threshold in absolute value are pruned,
            mask    =   abs(weight) >= threshold
        if sparsity level given, the smallest (sparsity)% weights in absolute value are pruned
        (or the largest (100-sparsity)% weights in absolute value are used)
            mask    =   topk(abs(weight), ret_typ='mask', k=weight.size*(100-sparsity)/100)

        => mask[i,j]    =   {0 if weight[i,j] is pruned, 1 otherwise} (for a matrix representation)

        weight  =   weight  *   mask
        grad    =   grad    *   mask
        state   =   state   *   mask

    This optimizer accepts the following parameters in addition to those accepted
    by :class:`.SGD`.

    Parameters
    ----------
    pruning_switch_epoch : list of ints, optional
        The epochs at which there is a change in sparsity level (should be in ascending order).

    weight_sparsity : list of floats, optional
        The sparsity on the weights required on each iteration of sparse training.

    bias_sparsity : list of floats, optional
        The sparsity on the biases required on each iteration of sparse training.

    weight_threshold : list of floats, optional
        The absolute value threshold on the weights required on each iteration of sparse training.

    bias_threshold : list of floats, optional
        The absolute value threshold on the biases required on each iteration of sparse training.

    batches_per_epoch : int, optional
        The number of batches in each epoch.
        (The ceiling integer value of number_of_examples / batch_size)
    """
    def __init__(self, pruning_switch_epoch, batches_per_epoch,
                 weight_sparsity=None, bias_sparsity=None,
                 weight_threshold=None, bias_threshold=None, **kwargs):
        super(SparseSGD, self).__init__(**kwargs)

        self.masks = []
        self.masks_updated = False
        self.epoch = 0
        self.pruning_switch_epoch = pruning_switch_epoch
        self.batches_per_epoch = batches_per_epoch

        # get weight and bias sparsity percentages
        self.weight_sparsity = weight_sparsity
        self.bias_sparsity = bias_sparsity
        if weight_sparsity is not None:
            assert len(weight_sparsity) == len(bias_sparsity), \
                'weight_sparsity and bias_sparsity should have same length'
            assert len(weight_sparsity) == len(pruning_switch_epoch), \
                'pruning_switch_epoch and weight_sparsity should have same length'

        # get weight and bias sparsity thresholds
        self.weight_threshold = weight_threshold
        self.bias_threshold = bias_threshold
        if weight_threshold is not None:
            assert len(weight_threshold) == len(bias_threshold), \
                'weight_threshold and bias_threshold should have same length'
            assert len(weight_threshold) == len(pruning_switch_epoch), \
                'pruning_switch_epoch and weight_sparsity_threshold should have same length'

        # either percentages or thresholds must be given
        assert weight_sparsity is not None or weight_threshold is not None,\
            'weight_sparsity or weight_sparsity_threshold should be given'

    def update_masks(self, index, weight):
        """Updates the masks for sparse training.

        Parameters
        ----------
        index : int
            The index for weight.
        weight : NDArray
            The weight matrix.

        Returns
        -------
        boolean
            If the masks were changed
        """
        # determine number of updates without actually updating the count
        if index not in self._index_update_count:
            num_update = self.begin_num_update
        else:
            num_update = self._index_update_count[index]
        num_update += 1
        num_update = max(num_update, self.num_update)

        # calculate epoch
        epoch = int((num_update - 1) / self.batches_per_epoch) + 1

        # determine if masks need to be updated, and get corresponding parameters
        if index == 0:
            self.masks_updated = True
        if self.epoch != epoch:
            self.epoch = epoch
            if epoch == 1:
                self.masks_updated = False
                if self.weight_sparsity is not None:
                    logging.info(log + 'bias-sparsity={}, weight-sparsity={}'.format(self.bias_sparsity[0], self.weight_sparsity[0]))
                else:
                    logging.info(log + 'bias-threshold={}, weight-threshold={}'.format(self.bias_threshold[0], self.weight_threshold[0]))
            if self.pruning_switch_epoch[0] + 1 == epoch:
                self.masks_updated = False
                self.pruning_switch_epoch.pop(0)
                if self.weight_sparsity is not None:
                    self.weight_sparsity.pop(0)
                    self.bias_sparsity.pop(0)
                    logging.info(log + 'bias-sparsity={}, weight-sparsity={}'.format(self.bias_sparsity[0], self.weight_sparsity[0]))
                else:
                    self.weight_threshold.pop(0)
                    self.bias_threshold.pop(0)
                    logging.info(log + 'bias-threshold={}, weight-threshold={}'.format(self.bias_threshold[0], self.weight_threshold[0]))

        # update masks if needed
        if not self.masks_updated:
            # initialize masks
            if epoch == 1:
                self.masks.append(None)
            # if percentages are given
            if self.weight_sparsity is not None:
                if len(weight.shape) == 1:
                    sparsity = self.bias_sparsity[0]
                else:
                    sparsity = self.weight_sparsity[0]
                number_unpruned = int((100.0 - sparsity) * weight.size / 100.0)
                self.masks[index] = topk(NDabs(weight), axis=None, ret_typ='mask',
                                         k=number_unpruned)
            # if thresholds are given
            else:
                if len(weight.shape) == 1:
                    threshold = self.bias_threshold[0]
                else:
                    threshold = self.weight_threshold[0]
                self.masks[index] = NDabs(weight) >= threshold

        return not self.masks_updated

    def update(self, index, weight, grad, state):
        assert(isinstance(weight, NDArray))
        assert(isinstance(grad, NDArray))

        # preprocessing for pruning
        if self.update_masks(index, weight):
            weight[:] = weight * self.masks[index]
        grad[:] = grad * self.masks[index]
        if state is not None:
            state[:] = state * self.masks[index]

        super(SparseSGD, self).update(index, weight, grad, state)