# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import tensorflow as tf

from nni.compression.tensorflow import Pruner

__all__ = [
    'LevelPruner',
    'SlimPruner',
]

class OneshotPruner(Pruner):
    def __init__(self, model, config_list, masker_class, **algo_kwargs):
        super().__init__(model, config_list)
        self.set_wrappers_attribute('calculated', False)
        self.masker = masker_class(model, self, **algo_kwargs)

    def validate_config(self, model, config_list):
        pass  # TODO

    def calc_masks(self, wrapper, wrapper_idx=None):
        if wrapper.calculated:
            return None
        sparsity = wrapper.config['sparsity']
        masks = self.masker.calc_masks(sparsity, wrapper, wrapper_idx)
        if masks is not None:
            wrapper.calculated = True
        return masks


class LevelPruner(OneshotPruner):
    def __init__(self, model, config_list):
        super().__init__(model, config_list, LevelPrunerMasker)


class SlimPruner(OneshotPruner):
    def __init__(self, model, config_list):
        super().__init__(model, config_list, SlimPrunerMasker)


class WeightMasker:
    def __init__(self, model, pruner, **kwargs):
        self.model = model
        self.pruner = pruner

    def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
        raise NotImplementedError()


class LevelPrunerMasker(WeightMasker):
    def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
        masks = {}
        for weight_variable in wrapper.layer.weights:
            if 'bias' in weight_variable.name:
                continue

            num_prune = int(tf.size(weight_variable).numpy() * sparsity)
            if num_prune == 0:
                continue

            weight = weight_variable.read_value()
            if wrapper.masks.get(weight_variable.name) is not None:
                weight = tf.math.multiply(weight, wrapper.masks[weight_variable.name])

            w_abs = tf.math.abs(weight)
            k = tf.size(weight) - num_prune
            topk = tf.math.top_k(tf.reshape(w_abs, [-1]), k).values
            if tf.size(topk) == 0:
                mask = tf.zeros_like(weight)
            else:
                mask = tf.math.greater_equal(w_abs, topk[-1])
            masks[weight_variable.name] = tf.cast(mask, weight.dtype)
        return masks

class SlimPrunerMasker(WeightMasker):
    def __init__(self, model, pruner, **kwargs):
        super().__init__(model, pruner)
        weight_list = []
        for wrapper in pruner.wrappers:
            weights = [w for w in wrapper.layer.weights if '/gamma:' in w.name]
            assert len(weights) == 1, f'Bad weights: {[w.name for w in wrapper.layer.weights]}'
            weight_list.append(tf.math.abs(weights[0].read_value()))
        all_bn_weights = tf.concat(weight_list, 0)
        k = int(all_bn_weights.shape[0] * pruner.wrappers[0].config['sparsity'])
        top_k = -tf.math.top_k(-tf.reshape(all_bn_weights, [-1]), k).values
        self.global_threshold = top_k.numpy()[-1]

    def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
        assert isinstance(wrapper.layer, tf.keras.layers.BatchNormalization), \
                'SlimPruner only supports 2D batch normalization layer pruning'

        weight = None
        weight_name = None
        bias_name = None

        for variable in wrapper.layer.weights:
            if '/gamma:' in variable.name:
                weight = variable.read_value()
                weight_name = variable.name
            elif '/beta:' in variable.name:
                bias_name = variable.name

        assert weight is not None
        if wrapper.masks.get(weight_name) is not None:
            weight *= wrapper.masks[weight_name]

        mask = tf.cast(tf.math.abs(weight) > self.global_threshold, weight.dtype)

        masks = {weight_name: mask}
        if bias_name:
            masks[bias_name] = mask
        return masks
