from typing import Any, Optional

import torch
from torch import nn
from torch.ao.quantization import QConfig


__all__ = ["QuantStub", "DeQuantStub", "QuantWrapper"]


class QuantStub(nn.Module):
    r"""Quantize stub module, before calibration, this is same as an observer,
    it will be swapped as `nnq.Quantize` in `convert`.

    Args:
        qconfig: quantization configuration for the tensor,
            if qconfig is not provided, we will get qconfig from parent modules
    """

    def __init__(self, qconfig: Optional[QConfig] = None):
        super().__init__()
        if qconfig:
            self.qconfig = qconfig

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x


class DeQuantStub(nn.Module):
    r"""Dequantize stub module, before calibration, this is same as identity,
    this will be swapped as `nnq.DeQuantize` in `convert`.

    Args:
        qconfig: quantization configuration for the tensor,
            if qconfig is not provided, we will get qconfig from parent modules
    """

    def __init__(self, qconfig: Optional[Any] = None):
        super().__init__()
        if qconfig:
            self.qconfig = qconfig

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x


class QuantWrapper(nn.Module):
    r"""A wrapper class that wraps the input module, adds QuantStub and
    DeQuantStub and surround the call to module with call to quant and dequant
    modules.

    This is used by the `quantization` utility functions to add the quant and
    dequant modules, before `convert` function `QuantStub` will just be observer,
    it observes the input tensor, after `convert`, `QuantStub`
    will be swapped to `nnq.Quantize` which does actual quantization. Similarly
    for `DeQuantStub`.
    """

    quant: QuantStub
    dequant: DeQuantStub
    module: nn.Module

    def __init__(self, module: nn.Module):
        super().__init__()
        qconfig = getattr(module, "qconfig", None)
        self.add_module("quant", QuantStub(qconfig))
        self.add_module("dequant", DeQuantStub(qconfig))
        self.add_module("module", module)
        self.train(module.training)

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        X = self.quant(X)
        X = self.module(X)
        return self.dequant(X)