# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
from torch import autograd
import math


class TopKBinarizer(autograd.Function):
    """
    Top-k Binarizer.
    Computes a binary mask M from a real value matrix S such that `M_{i,j} = 1` if and only if `S_{i,j}`
    is among the k% highest values of S.
    Implementation is inspired from:
        https://github.com/yaozhewei/MLPruning
    """

    @staticmethod
    def forward(ctx, inputs: torch.tensor, threshold: float, sigmoid: bool):
        """
        Args:
            inputs (`torch.FloatTensor`)
                The input matrix from which the binarizer computes the binary mask.
            threshold (`float`)
                The percentage of weights to keep (the rest is pruned).
                `threshold` is a float between 0 and 1.
            sigmoid (`bool`)
                Whether to apply a sigmoid on the threshold
        Returns:
            mask (`torch.FloatTensor`)
                Binary matrix of the same size as `inputs` acting as a mask (1 - the associated weight is
                retained, 0 - the associated weight is pruned).
        """
        # Get the subnetwork by sorting the inputs and using the top threshold
        if sigmoid:
            threshold = torch.sigmoid(threshold).item()
        ctx.sigmoid = sigmoid
        mask = inputs.clone()

        _, idx = inputs.flatten().sort(descending=True)
        j = math.ceil(threshold * inputs.numel())

        # flat_out and mask access the same memory.
        flat_out = mask.flatten()
        flat_out[idx[j:]] = 0.
        flat_out[idx[:j]] = 1.
        ctx.save_for_backward(mask)

        return mask

    @staticmethod
    def backward(ctx, gradOutput):
        mask, = ctx.saved_tensors
        if ctx.sigmoid:
            return gradOutput.clone(), ((gradOutput * mask).sum()).view(-1), None
        else:
            return gradOutput.clone(), None, None


class SymQuantizer(torch.autograd.Function):
    """
    Symmetric quantization
    """

    @staticmethod
    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
        """
        Args:
            inputs (`torch.FloatTensor`)
                The input which needs to be quantized
            num_bits (int, >=4)
                Number of bits to use for quantization
            min_value/max_value (torch.FloatTensor)
                Used for static activation quantization
            num_groups (int)
                How many groups to partition the quantization into
        Returns:
            quantized_input (`torch.FloatTensor`)
                Quantized input
        """
        assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None
                                                             and num_groups == 1)
        q_range = 2**num_bits
        input_shape = input.shape
        if min_value is None:
            input = input.reshape(num_groups, -1)
            max_input = torch.amax(torch.abs(input), dim=-1).view(num_groups, -1)
        else:
            max_input = torch.max(min_value.abs(), max_value).view(-1)

        scale = 2 * max_input / q_range
        output = (input / scale).round().clamp(-q_range // 2, q_range // 2 - 1) * scale
        output = output.reshape(input_shape).contiguous()
        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()
        return grad_input, None, None, None, None


class AsymQuantizer(torch.autograd.Function):
    """
    Asymmetric quantization
    """

    @staticmethod
    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
        """
        Args:
            inputs (`torch.FloatTensor`)
                The input which needs to be quantized
            num_bits (int, >=4)
                Number of bits to use for quantization
            min_value/max_value (torch.FloatTensor)
                Used for static activation quantization
            num_groups (int)
                How many groups to partition the quantization into
        Returns:
            quantized_input (`torch.FloatTensor`)
                Quantized input
        """

        assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None
                                                             and num_groups == 1)
        q_range = 2**num_bits
        input_shape = input.shape
        if min_value is None:
            input = input.reshape(num_groups, -1)
            min_value = input.amin(dim=-1, keepdim=True)
            max_value = input.amax(dim=-1, keepdim=True)

        scale = (max_value - min_value) / q_range
        zero_point = (min_value / scale).round() * scale

        output = ((input - zero_point) / scale).round().clamp(0, q_range - 1) * scale + zero_point
        output = output.reshape(input_shape).contiguous()
        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()
        return grad_input, None, None, None, None


class TernaryQuantizer(torch.autograd.Function):
    """
    Ternary quantization
    """

    @staticmethod
    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
        """
        Args:
            inputs (`torch.FloatTensor`)
                The input which needs to be quantized
            num_bits (int)
                Dummy variable
            min_value/max_value (torch.FloatTensor)
                Used for static activation quantization; for now they are dummy variable
            num_groups (int)
                How many groups to partition the quantization into
        Returns:
            quantized_input (`torch.FloatTensor`)
                Quantized input
        """

        assert (min_value is None and max_value is None)
        input_flat = input.reshape(num_groups, -1)
        n = input_flat.shape[1]
        m = input_flat.norm(p=1, dim=1).div(n)
        thres = (0.7 * m).view(-1, 1)
        pos = (input_flat > thres).type(input.type())
        neg = (input_flat < -thres).type(input.type())
        mask = (input_flat.abs() > thres).type(input.type())
        alpha = ((mask * input_flat).abs().sum(dim=1) / mask.sum(dim=1)).view(-1, 1)
        output = alpha * pos - alpha * neg
        output = output.reshape(input.shape).contiguous()
        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()
        return grad_input, None, None, None, None


class BinaryQuantizer(torch.autograd.Function):
    """
    Binary quantization
    """

    @staticmethod
    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
        """
        Args:
            inputs (`torch.FloatTensor`)
                The input which needs to be quantized
            num_bits (int)
                Dummy variable
            min_value/max_value (torch.FloatTensor)
                Used for static activation quantization; for now they are dummy variable
            num_groups (int)
                How many groups to partition the quantization into
        Returns:
            quantized_input (`torch.FloatTensor`)
                Quantized input
        """

        assert (min_value is None and max_value is None)
        input_flat = input.reshape(num_groups, -1)
        n = input_flat.shape[1]
        m = input_flat.norm(p=1, dim=1, keepdim=True).div(n)
        output = input_flat.sign().mul(m)
        output = output.reshape(input.shape).contiguous()
        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()
        return grad_input, None, None, None, None