# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
# Copyright (c) 2021 NVIDIA Corporation.  All rights reserved.
# 
# Licensed 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.
"""
Utilities of Auto SParsity (ASP).
"""

from __future__ import print_function

import sys
import math
import collections
import numpy as np
from enum import Enum
from itertools import permutations
import threading

__all__ = [
    'calculate_density', 'check_mask_1d', 'get_mask_1d', 'check_mask_2d',
    'get_mask_2d_greedy', 'get_mask_2d_best', 'create_mask', 'check_sparsity',
    'MaskAlgo', 'CheckMethod'
]


class MaskAlgo(Enum):
    r"""
    A collection of all mask generating algorithms.
    There currently are three algorithms, `MASK_1D`, `MASK_2D_GREEDY` and `MASK_2D_BEST`
    """
    MASK_1D = 'get_mask_1d'
    MASK_2D_GREEDY = 'get_mask_2d_greedy'
    MASK_2D_BEST = 'get_mask_2d_best'


class CheckMethod(Enum):
    r"""
    A collection of all sparsity checking approaches.
    There currently are two methods, `CHECK_1D` and `CHECK_2D`
    """
    CHECK_1D = 'check_mask_1d'
    CHECK_2D = 'check_mask_2d'

    @staticmethod
    def get_checking_method(mask_algo):
        r"""
        Get sparsity checking method by mask generating algorithm.

        Args:
            mask_algo (MaskAlgo): The algorithm of mask generating.
        Returns:
            CheckMethod: The corresponded sparsity checking method.
        Examples:
            .. code-block:: python

            import numpy as np
            from paddle.fluid.contrib.sparsity import MaskAlgo, CheckMethod

            CheckMethod.get_checking_method(MaskAlgo.MASK_1D)
            # CheckMethod.CHECK_1D

            CheckMethod.get_checking_method(MaskAlgo.MASK_2D_GREEDY)
            # CheckMethod.CHECK_2D

            CheckMethod.get_checking_method(MaskAlgo.MASK_2D_BEST)
            # CheckMethod.CHECK_2D
        """
        assert isinstance(mask_algo, MaskAlgo), \
               "mask_algo should be MaskAlgo type"
        if mask_algo == MaskAlgo.MASK_1D:
            return CheckMethod.CHECK_1D
        else:
            return CheckMethod.CHECK_2D


def calculate_density(x):
    r"""
    Return the density of the input tensor.

    Args:
        x (nparray): The input tensor.
    Returns:
        float: The density of :attr:`x`.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          x = np.array([[0, 1, 3, 0],
                        [1, 1, 0, 1]])
          sparsity.calculate_density(x) # 0.625
    """
    x_flattened = x.flatten()
    return float(np.nonzero(x_flattened)[0].size) / x_flattened.size


def _reshape_1d(mat, m):
    r"""
    Reshape the input 2D matrix to shape (-1, m).
    If the second dimension of :attr:`mat` is not a multiples of :attr:`m`, 
    then this function would pad the remainder with 0 before reshaping.

    .. math::

        remainder = mat.shape[1] % m

    Args:
        mat (nparray): The input 2D matrix.
        m (int): The second dimension of reshaped matrix.
    Returns:
        tuple: A pair of the reshaped and padded matrix and the shape of padded matrix (non-reshaping).
    """
    assert len(mat.shape) == 2, "The input mat should be a 2D matrix!"

    remainder = mat.shape[1] % m
    if mat.shape[1] % m > 0:
        mat_padded = np.zeros((mat.shape[0], mat.shape[1] + (m - remainder)))
        mat_padded[:, :mat.shape[1]] = mat
        shape = mat_padded.shape
        return mat_padded.reshape(-1, m), shape
    else:
        return mat.reshape(-1, m), mat.shape


def check_mask_1d(mat, n, m):
    r"""
    Check if every row of the input matrix :attr:`mat` is in 1D `n:m` sparse pattern.
    This function would pad the second dimension of :attr:`mat` by zero 
    to be a multiples of :attr:`m` if necessary.

    1D `n:m` sparse pattern: At least :attr:`n` zeros in every :math:`1 \times m` block.

    Args:
        mat (nparray): The input matrix.
        n (int): n of `n:m` sparse pattern.
        m (int): m of `n:m` sparse pattern.
    Returns:
        bool: True if every row of :attr:`mat` is in 1D n:m sparse pattern, else False.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          x = np.array([[0, 1, 3, 0],
                        [1, 0, 0, 1]])
          sparsity.check_mask_1d(x, 2, 4) # True

          x = np.array([[0, 1, 5, 4],
                        [1, 0, 0, 1]])
          sparsity.check_mask_1d(x, 2, 4) # False

          # x would be padded to shape (2, 8)
          x = np.array([[0, 1, 0, 4, 6],
                        [1, 0, 0, 1, 7]])
          sparsity.check_mask_1d(x, 2, 4) # True
    """
    if len(mat.shape) <= 1:
        mat_flattern, shape = _reshape_1d(mat.reshape(1, mat.shape[0]), m)
    else:
        mat_flattern, shape = _reshape_1d(mat, m)

    for sub_mat in mat_flattern:
        if np.nonzero(sub_mat)[0].size > (m - n):
            return False
    return True


def get_mask_1d(mat, n, m):
    r"""
    Generate 1D `n:m` sparse pattern mask of the input matrix :attr:`mat` 
    in row-directory. This function would pad the second dimension of :attr:`mat` 
    by zero to be a multiples of :attr:`m` before mask generation.

    1D `n:m` sparse pattern: At least :attr:`n` zeros in every :math:`1 \times m` block.

    Args:
        mat (nparray): The input matrix.
        n (int): n of `n:m` sparse pattern.
        m (int): m of `n:m` sparse pattern.
    Returns:
        nparray: The 1D `n:m` sparse mask of :attr:`mat`.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          mat = np.array([[0, 1, 5, 4],
                          [2, 7, 3, 6]])
          mask = sparsity.get_mask_1d(mat, 2, 4)
          # nparray([[0, 0, 1, 1],
          #          [0, 1, 0, 1]])
          sparsity.check_mask_1d(mask, 2, 4) # True
    """
    mat_flattern, shape = _reshape_1d(mat, m)

    mask_flattern = np.ones_like(mat_flattern)
    mask = np.ones_like(mat)
    for i in range(mat_flattern.shape[0]):
        sub_mat = mat_flattern[i]
        min_order_indices = np.argsort(np.absolute(sub_mat))
        mask_flattern[i, min_order_indices[:n].tolist()] = 0
    mask_flattern = mask_flattern.reshape(shape)
    mask[:, :] = mask_flattern[:, :mat.shape[1]]
    return mask


def _reshape_2d(mat, m):
    r"""
    Reshape the input 2D matrix to shape (-1, :math:`m \times m`).
    In each dimension of :attr:`mat`, if it is not a multiples of :attr:`m`, 
    then this function would pad the remainder with 0 before reshaping.

    .. math::

        remainder_0 = mat.shape[0] % m \\
        remainder_1 = mat.shape[1] % m

    Args:
        mat (nparray): The input 2D matrix.
        m (int): The square root of second dimension of reshaped matrix.
    Returns:
        tuple: A pair of the reshaped and padded matrix and the shape of padded matrix (non-reshaping).
    """
    assert len(mat.shape) == 2, "The input mat should be a 2D matrix!"

    remainder_0 = mat.shape[0] % m
    remainder_1 = mat.shape[1] % m

    new_shape = (mat.shape[0] if remainder_0 == 0 \
                 else mat.shape[0] + (m - remainder_0),
                 mat.shape[1] if remainder_1 == 0 \
                 else mat.shape[1] + (m - remainder_1))
    mat_padded = np.zeros(new_shape)
    mat_padded[:mat.shape[0], :mat.shape[1]] = mat

    mat_flattern = np.empty(new_shape).reshape(-1, m * m)
    curr_idx = 0
    for row_start in range(0, mat_padded.shape[0], m):
        row_end = row_start + m
        for col_start in range(0, mat_padded.shape[1], m):
            col_end = col_start + m
            sub_mat = np.squeeze(mat_padded[row_start:row_end, \
                                            col_start:col_end] \
                                            .reshape(-1))
            mat_flattern[curr_idx] = sub_mat
            curr_idx += 1
    return mat_flattern, mat_padded.shape


def check_mask_2d(mat, n, m):
    r"""
    Check if every :math:`m \times m` block of the input matrix :attr:`mat` is in 2D `n:m` sparse pattern.
    This function would pad each dimension of :attr:`mat` by zero to be a multiples of 
    :attr:`m` if necessary.

    2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block 
    under the constraint of at least :attr:`n` zeros for each row and column.

    Args:
        mat (nparray): The input matrix.
        n (int): n of `n:m` sparse pattern.
        m (int): m of `n:m` sparse pattern.
    Returns:
        bool: True if  every :math:`m \times m` block of the input matrix :attr:`mat` is in 2D `n:m` sparse pattern, else False.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          x = np.array([[0, 8, 9, 0],
                        [9, 0, 0, 10],
                        [5, 0, 0, 6],
                        [0, 4, 6, 0]])
          sparsity.check_mask_2d(x, 2, 4) # True

          x = np.array([[0, 8, 0, 9],
                        [9, 0, 0, 10],
                        [0, 5, 0, 6],
                        [0, 4, 6, 0]])
          sparsity.check_mask_2d(x, 2, 4) # False

          # x would be padded to shape (8, 8)
          x = np.array([[0, 8, 0, 9],
                        [9, 0, 7, 0],
                        [0, 5, 0, 6],
                        [3, 0, 6, 0],
                        [1, 1, 0, 1]])
          sparsity.check_mask_2d(x, 2, 4) # True
    """
    mat_padded, shape = _reshape_2d(mat, m)
    for sub_mat in mat_padded:
        sub_mask = np.absolute(np.squeeze(sub_mat.reshape(m, m))) > 0
        if (np.sum(np.sum(sub_mask, axis=1) > (m-n)) != 0) and \
            (np.sum(np.sum(sub_mask, axis=0) > (m-n)) != 0):
            return False
    return True


def get_mask_2d_greedy(mat, n, m):
    r"""
    Greedily generate 2D `n:m` sparse pattern mask of the input matrix :attr:`mat`. 
    This function would pad each dimension of :attr:`mat` by zero to be a multiples of :attr:`m` before mask generation.

    2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block 
    under the constraint of at least :attr:`n` zeros for each row and column.
    Greedily generating: For each :math:`m \times m` block, selecting values to keep in descent order.

    Args:
        mat (nparray): The input matrix.
        n (int): n of `n:m` sparse pattern.
        m (int): m of `n:m` sparse pattern.
    Returns:
        nparray: The 2D `n:m` sparse mask of :attr:`mat`.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          mat = np.array([[9, 8, 3, 7],
                          [9, 2, 1, 10],
                          [5, 1, 3, 6],
                          [2, 4, 6, 1]])
          mask = sparsity.get_mask_2d_greedy(mat, 2, 4)
          # nparray([[1. 1. 0. 0.]
          #          [1. 0. 0. 1.]
          #          [0. 0. 1. 1.]
          #          [0. 1. 1. 0.]])
          sparsity.check_mask_2d(mask, 2, 4) # True
    """
    mat_padded, shape = _reshape_2d(mat, m)
    mask_padded = np.zeros_like(mat_padded).reshape(-1, m, m)

    for idx in range(len(mat_padded)):
        sub_mat = np.absolute(np.squeeze(mat_padded[idx]))
        sub_mask = np.squeeze(mask_padded[idx])

        min_order_1d_indices = np.argsort(sub_mat)
        min_order_2d_indices = [(int(x / m), x % m)
                                for x in min_order_1d_indices]
        row_counter = collections.Counter()
        col_counter = collections.Counter()

        for i in range(len(min_order_1d_indices) - 1, -1, -1):
            matrix_entry = min_order_2d_indices[i]
            if (row_counter[matrix_entry[0]] == n) or \
               (col_counter[matrix_entry[1]] == n):
                continue

            sub_mask[matrix_entry[0], matrix_entry[1]] = 1.0
            row_counter[matrix_entry[0]] += 1
            col_counter[matrix_entry[1]] += 1

    mask = np.empty(shape)
    curr_idx = 0
    for row_start in range(0, shape[0], m):
        row_end = row_start + m
        for col_start in range(0, shape[1], m):
            col_end = col_start + m
            mask[row_start:row_end, col_start:col_end] = mask_padded[curr_idx]
            curr_idx += 1
    return mask[:mat.shape[0], :mat.shape[1]]


_valid_2d_patterns_lock = threading.Lock()
_valid_2d_patterns = {}


def _compute_valid_2d_patterns(n, m):
    r"""
    Compute all vaild 2D `n:m` sparse patterns.

    2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block 
    under the constraint of at least :attr:`n` zeros for each row and column.

    Args:
        n (int): n of `n:m` sparse pattern.
        m (int): m of `n:m` sparse pattern.
    Returns:
        dictionary: A dictionary with key: *m_n* (string) and value: all vaild 2D `n:m` sparse patterns.
    """
    global _valid_2d_patterns_lock
    global _valid_2d_patterns

    valid_key = '{}_{}'.format(m, n)
    if valid_key in _valid_2d_patterns:
        return _valid_2d_patterns[valid_key]
    else:
        patterns = np.zeros(m)
        patterns[:n] = 1
        patterns = list(set(permutations(patterns.tolist())))
        patterns = patterns + patterns
        patterns = np.asarray(list(set(permutations(patterns, m))))

        valid = ((patterns.sum(axis=1) <= n).sum(axis=1) == m
                 ).nonzero()[0].reshape(-1)
        valid_patterns = np.empty((valid.shape[0], m, m))
        valid_patterns[:] = patterns[valid[:]]

        _valid_2d_patterns_lock.acquire()
        _valid_2d_patterns[valid_key] = valid_patterns
        _valid_2d_patterns_lock.release()

        return valid_patterns


def get_mask_2d_best(mat, n, m):
    r"""
    Generate 2D `n:m` sparse pattern mask of the input matrix :attr:`mat` 
    to form sparse matrix with maximun L1 norm .This function would pad each 
    dimension of :attr:`mat` by zero to be a multiples of :attr:`m` before mask generation.

    2D `n:m` sparse pattern: At least :math:`n \times n` zeros in every :math:`m \times m` block 
    under the constraint of at least :attr:`n` zeros for each row and column.

    *Note*: L1 norm of sparse matrix from `Best` API is greater than or equal to the one from `Greedy`.

    Args:
        mat (nparray): The input matrix.
        n (int): n of `n:m` sparse pattern.
        m (int): m of `n:m` sparse pattern.
    Returns:
        nparray: The 1D `n:m` sparse mask of :attr:`mat`.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          mat = np.array([[2, 8, 9, 9],
                          [9, 1, 3, 9],
                          [5, 6, 3, 9],
                          [2, 4, 6, 9]])
          mask_greedy = sparsity.get_mask_2d_greedy(mat, 2, 4)
          mask_greedy = sparsity.get_mask_2d_best(mat, 2, 4)
          print("L1 norm of `greedy` sparse matrix", np.multiply(mat, mask_greedy).sum()) # 56
          print("L1 norm of `best` sparse matrix", np.multiply(mat, mask_best).sum()) # 61
    """
    patterns = _compute_valid_2d_patterns(n, m)

    mat_flattern, shape = _reshape_2d(mat, m)
    mask_flattern = np.ones_like(mat_flattern).reshape(-1, m, m)
    pmax = np.argmax(
        np.matmul(mat_flattern, patterns.reshape(patterns.shape[0], m * m).T),
        axis=1)

    mask_flattern[:] = patterns[pmax[:]]
    mask = np.empty(shape)

    curr_idx = 0
    for row_start in range(0, shape[0], m):
        row_end = row_start + m
        for col_start in range(0, shape[1], m):
            col_end = col_start + m
            mask[row_start:row_end, col_start:col_end] = mask_flattern[curr_idx]
            curr_idx += 1
    return mask[:mat.shape[0], :mat.shape[1]]


def create_mask(tensor, func_name=MaskAlgo.MASK_1D, n=2, m=4):
    r"""
    Create `n:m` sparse pattern mask of the input tensor via function given by :attr:`func_name`.
    Currently only support tensor with dimension less than or equal to 4.

    Args:
        tensor (nparray): The input tensor.
        func_name (MaskAlgo, optional): The function name to generate spase mask. Default is `MaskAlgo.MASK_1D`. All options please refer to `MaskAlgo`.
        n (int, optional): n of `n:m` sparse pattern. Default is 2.
        m (int, optional): m of `n:m` sparse pattern. Default is 4.
    Returns:
        nparray: The `n:m` sparse mask of :attr:`tensor` generated by :attr:`func_name`.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          tensor = np.array([[2, 8, 9, 9],
                             [9, 1, 3, 9],
                             [5, 6, 3, 9],
                             [2, 4, 6, 9]])
          mask_1d = sparsity.create_mask(tensor, func_name=sparsity.MaskAlgo.MASK_1D)
          # nparray([[0 0 1 1],
          #          [1 0 0 1],
          #          [0 1 0 1],
          #          [0 0 1 1]])
          mask_2d = sparsity.create_mask(tensor, func_name=sparsity.MaskAlgo.MASK_2D_BEST)
          # nparray([[0 1 1 0],
          #          [1 0 0 1],
          #          [1 1 0 0],
          #          [0 0 1 1]])
    """
    shape = tensor.shape
    dtype = tensor.dtype
    t = tensor.astype(float)

    assert isinstance(func_name, MaskAlgo), \
           "func_name argumet of create_mask is only accepted as type MaskAlgo. " \
           "But got {}".format(type(func_name))
    func = getattr(sys.modules[__name__], func_name.value, None)
    if len(shape) == 1:
        t = t.reshape(1, shape[0])
    elif len(shape) == 2:
        t = t.reshape(shape[0], shape[1])
    elif len(shape) == 3:
        t = t.reshape(shape[0] * shape[1], shape[2])
    # 4d-tensor conv (h, w, in, out) -> (h*w*out, in) in GemmConvKernel Op
    elif len(shape) == 4:
        t = t.transpose([0, 1, 3, 2]).reshape(shape[0] * shape[1] * shape[3],
                                              shape[2])
        mask = func(t, n=n, m=m)
        return mask.reshape([shape[0], shape[1], shape[3],
                             shape[2]]).transpose([0, 1, 3, 2]).astype(dtype)
    else:
        raise ValueError("The dimension of input tensor is not supported in create_mask, " \
                         "Only dimension < 4 is supported but got {}".format(len(shape)))

    mask = func(t, n=n, m=m)
    return mask.reshape(shape).astype(dtype)


def check_sparsity(tensor, func_name=CheckMethod.CHECK_1D, n=2, m=4):
    r"""
    Check if input tensor is in `n:m` sparse pattern via function given by :attr:`func_name`.
    Currently only support tensor with dimension less than or equal to 4.

    Args:
        tensor (nparray): The input tensor.
        func_name (CheckMethod, optional): The function name to generate spase mask. Default is `CheckMethod.CHECK_1D`. All options please refer to `CheckMethod`.
        n (int, optional): n of `n:m` sparse pattern. Default is 2.
        m (int, optional): m of `n:m` sparse pattern. Default is 4.
    Returns:
        bool: True if tensor pass checking of function given by :attr:`func_name`, else False.
    Examples:
        .. code-block:: python

          import numpy as np
          import paddle.fluid.contrib.sparsity as sparsity

          tensor = np.array([[2, 8, 9, 9],
                             [9, 1, 3, 9],
                             [5, 6, 3, 9],
                             [2, 4, 6, 9]])
          mask_1d = sparsity.create_mask(tensor, func_name=sparsity.MaskAlgo.MASK_1D)
          # nparray([[0 0 1 1],
          #          [1 0 0 1],
          #          [0 1 0 1],
          #          [0 0 1 1]])
          sparsity.check_sparsity(mask_1d, func_name=sparsity.CheckMethod.CHECK_1D) # True
          sparsity.check_sparsity(mask_1d, func_name=sparsity.CheckMethod.CHECK_2D) # False
    """
    shape = tensor.shape
    t = tensor.astype(float)

    assert type(func_name) == CheckMethod, \
           "func_name argumet of check_sparsity is only accepted as type CheckMethod. " \
           "But got {}".format(type(func_name))
    func = getattr(sys.modules[__name__], func_name.value, None)
    if len(shape) == 1:
        t = t.reshape(1, shape[0])
    elif len(shape) == 2:
        t = t.reshape(shape[0], shape[1])
    elif len(shape) == 3:
        t = t.reshape(shape[0] * shape[1], shape[2])
    # 4d-tensor conv (h, w, in, out) -> (h*w*out, in) in GemmConvKernel Op
    elif len(shape) == 4:
        t = t.transpose([0, 1, 3, 2]).reshape(
            [shape[0] * shape[1] * shape[3], shape[2]])
    else:
        raise ValueError("The dimension of input tensor is not supported in create_mask, " \
                         "Only dimension < 4 is supported but got {}".format(len(shape)))

    return func(t, n=n, m=m)