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

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.static.sparsity import MaskAlgo
            from paddle.fluid.contrib.sparsity import 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 paddle
            import numpy as np

            x = np.array([[0, 1, 3, 0],
                        [1, 1, 0, 1]])
            paddle.incubate.asp.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_best = 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)