# mypy: allow-untyped-defs
from typing import Optional

import torch
import torch.ao.nn.intrinsic as nni
from torch.ao.nn.quantized.modules.utils import (
    _hide_packed_params_repr,
    _quantize_weight,
)
from torch.ao.nn.sparse.quantized import linear
from torch.ao.nn.sparse.quantized.utils import LinearBlockSparsePattern


__all__ = ["Linear"]


class Linear(torch.nn.Module):
    r"""
    A dynamically quantized sparse linear module with float tensor as inputs and outputs.
    """

    _version = 1
    _op_type = "sparse_dynamic"
    _FLOAT_MODULE = torch.nn.Linear

    def __init__(
        self,
        in_features,
        out_features,
        row_block_size,
        col_block_size,
        bias=True,
        dtype=torch.qint8,
    ):
        super().__init__()

        if dtype != torch.qint8:
            raise NotImplementedError(
                "Only QINT8 is supported for Sparse Quantized Linear Dynamic"
            )

        self.in_features = in_features
        self.out_features = out_features

        if bias:
            bias = torch.zeros(self.out_features, dtype=torch.float)
        else:
            bias = None

        qweight = torch._empty_affine_quantized(
            [out_features, in_features], scale=1, zero_point=0, dtype=torch.qint8
        )
        self._packed_params = linear.LinearPackedParams(
            row_block_size=row_block_size, col_block_size=col_block_size, dtype=dtype
        )
        self._packed_params.set_weight_bias(
            qweight, bias, row_block_size, col_block_size
        )

    def _get_name(self):
        return "SparseQuantizedDynamicLinear"

    def extra_repr(self):
        return f"in_features={self.in_features}, out_features={self.out_features}, qscheme={self.weight().qscheme()}"

    def __repr__(self):
        return _hide_packed_params_repr(self, linear.LinearPackedParams)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.ops.sparse.qlinear_dynamic(x, self._packed_params._packed_params)

    def _save_to_state_dict(self, destination, prefix, keep_vars):
        super()._save_to_state_dict(destination, prefix, keep_vars)
        destination[prefix + "op_type"] = self._op_type

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        op_type = int(state_dict[prefix + "op_type"])
        assert op_type == "sparse", (
            f"Cannot load from op_type [{op_type}], expecting [{self._op_type}]"
        )
        state_dict.pop(prefix + "op_type")

        version = local_metadata.get("version", None)
        assert version <= self._version

        # Is this code valid? In old quantization it seemed to be used to load
        # older model
        weight = state_dict.pop(prefix + "weight")
        bias = state_dict.pop(prefix + "bias")
        state_dict.update(
            {
                prefix + "_packed_params.weight": weight,
                prefix + "_packed_params.bias": bias,
            }
        )

        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            False,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def _weight_bias(self):
        return self._packed_params._weight_bias()

    def weight(self):
        return self._weight_bias()[0]

    def bias(self):
        return self._weight_bias()[1]

    def set_weight_bias(
        self,
        w: torch.Tensor,
        b: Optional[torch.Tensor],
        row_block_size: Optional[int],
        col_block_size: Optional[int],
    ) -> None:
        assert row_block_size is not None and col_block_size is not None
        self.out_features = w.shape[0]
        self.in_features = w.shape[1]
        self._packed_params.set_weight_bias(w, b, row_block_size, col_block_size)

    @classmethod
    def from_float(cls, mod, use_precomputed_fake_quant=False):
        r"""Create a quantized sparse dynamic module from a float module.

        We only care about the convert at this stage, no need for observers just yet.
        """
        assert type(mod) == cls._FLOAT_MODULE, (
            " nnq."
            + cls.__name__
            + ".from_float only works for "
            + cls._FLOAT_MODULE.__name__
        )
        # TODO: Need to add options to qconfig to avoid the calibration.
        # TODO: Add calibration for the sparsity
        assert hasattr(mod, "qconfig"), "Input float module must have qconfig defined"
        if type(mod) == nni.LinearReLU:
            mod = mod[0]
        if mod.qconfig is not None and mod.qconfig.weight is not None:
            weight_observer = mod.qconfig.weight()
        else:
            # We have the circular import issues if we import the qconfig in the beginning of this file:
            # https://github.com/pytorch/pytorch/pull/24231. The current workaround is to postpone the
            # import until we need it.
            from torch.ao.quantization.qconfig import default_dynamic_qconfig

            weight_observer = default_dynamic_qconfig.weight()

        # It is important to multiply by the mask BEFORE calling the `weight_observer`
        # TODO (zaf): Mask might not be part of the qconfig (T83295194)
        weight = mod.weight
        if getattr(mod.qconfig, "mask", False):
            weight = mod.qconfig.mask * mod.weight

        weight_observer(weight)
        dtype = weight_observer.dtype
        assert dtype == torch.qint8, "Weight observer must have dtype torch.qint8"
        _w_sc, w_zp = weight_observer.calculate_qparams()
        if isinstance(w_zp, torch.Tensor):
            assert not torch.any(w_zp.bool()), "All weight zero points must map to 0"
        else:
            assert w_zp == 0, "Weight zero point must map to 0"
        qweight = _quantize_weight(weight.float(), weight_observer)

        row_block_size, col_block_size = LinearBlockSparsePattern.block_size()
        qlinear = cls(
            mod.in_features,
            mod.out_features,
            row_block_size,
            col_block_size,
            dtype=dtype,
        )
        qlinear.set_weight_bias(qweight, mod.bias, row_block_size, col_block_size)
        return qlinear