# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from .em import EM, EmptyClusterResolveError


class PQ(EM):
    """
    Quantizes the layer weights W with the standard Product Quantization
    technique. This learns a codebook of codewords or centroids of size
    block_size from W. For further reference on using PQ to quantize
    neural networks, see "And the Bit Goes Down: Revisiting the Quantization
    of Neural Networks", Stock et al., ICLR 2020.

    PQ is performed in two steps:
    (1) The matrix W (weights or fully-connected or convolutional layer)
        is reshaped to (block_size, -1).
            - If W is fully-connected (2D), its columns are split into
              blocks of size block_size.
            - If W is convolutional (4D), its filters are split along the
              spatial dimension.
    (2) We apply the standard EM/k-means algorithm to the resulting reshaped matrix.

    Args:
        - W: weight matrix to quantize of size (in_features x out_features)
        - block_size: size of the blocks (subvectors)
        - n_centroids: number of centroids
        - n_iter: number of k-means iterations
        - eps: for cluster reassignment when an empty cluster is found
        - max_tentatives for cluster reassignment when an empty cluster is found
        - verbose: print information after each iteration

    Remarks:
        - block_size be compatible with the shape of W
    """

    def __init__(
        self,
        W,
        block_size,
        n_centroids=256,
        n_iter=20,
        eps=1e-6,
        max_tentatives=30,
        verbose=True,
    ):
        self.block_size = block_size
        W_reshaped = self._reshape(W)
        super(PQ, self).__init__(
            W_reshaped,
            n_centroids=n_centroids,
            n_iter=n_iter,
            eps=eps,
            max_tentatives=max_tentatives,
            verbose=verbose,
        )

    def _reshape(self, W):
        """
        Reshapes the matrix W as expained in step (1).
        """

        # fully connected: by convention the weight has size out_features x in_features
        if len(W.size()) == 2:
            self.out_features, self.in_features = W.size()
            assert (
                self.in_features % self.block_size == 0
            ), "Linear: n_blocks must be a multiple of in_features"
            return (
                W.reshape(self.out_features, -1, self.block_size)
                .permute(2, 1, 0)
                .flatten(1, 2)
            )

        # convolutional: we reshape along the spatial dimension
        elif len(W.size()) == 4:
            self.out_channels, self.in_channels, self.k_h, self.k_w = W.size()
            assert (
                self.in_channels * self.k_h * self.k_w
            ) % self.block_size == 0, (
                "Conv2d: n_blocks must be a multiple of in_channels * k_h * k_w"
            )
            return (
                W.reshape(self.out_channels, -1, self.block_size)
                .permute(2, 1, 0)
                .flatten(1, 2)
            )
        # not implemented
        else:
            raise NotImplementedError(W.size())

    def encode(self):
        """
        Performs self.n_iter EM steps.
        """

        self.initialize_centroids()
        for i in range(self.n_iter):
            try:
                self.step(i)
            except EmptyClusterResolveError:
                break

    def decode(self):
        """
        Returns the encoded full weight matrix. Must be called after
        the encode function.
        """

        # fully connected case
        if "k_h" not in self.__dict__:
            return (
                self.centroids[self.assignments]
                .reshape(-1, self.out_features, self.block_size)
                .permute(1, 0, 2)
                .flatten(1, 2)
            )

        # convolutional case
        else:
            return (
                self.centroids[self.assignments]
                .reshape(-1, self.out_channels, self.block_size)
                .permute(1, 0, 2)
                .reshape(self.out_channels, self.in_channels, self.k_h, self.k_w)
            )