import torch
from .quantize_base import QuantizeBase


class PseudoFakeQuantize(QuantizeBase):
    """This is actually torch.quantization.FakeQuantize.
    """
    def __init__(self, observer, **observer_kwargs):
        super(PseudoFakeQuantize, self).__init__(observer,  **observer_kwargs)
        self.register_buffer('scale', torch.tensor([1.0], dtype=torch.float))
        self.register_buffer('zero_point', torch.tensor([0], dtype=torch.int))


    def _forward_quant(self, X):
        if self.training:
            return X
        else:
            if self.fake_quant_enabled[0] == 1:
                if self.is_per_channel:
                    X = torch.fake_quantize_per_channel_affine(
                        X, self.scale.data, self.zero_point.data.long(),
                        self.ch_axis, self.quant_min, self.quant_max)
                else:
                    X = torch.fake_quantize_per_tensor_affine(
                        X, self.scale.item(), self.zero_point.item(),
                        self.quant_min, self.quant_max)
            return X

    @torch.jit.export
    def extra_repr(self):
        return 'fake_quant_enabled={}, observer_enabled={}, is_initialized_params={}, ' \
               'quant_min={}, quant_max={}, dtype={}, qscheme={}, ch_axis={}, ' \
               'scale={}, zero_point={} (test/eval)'.format(
                   self.fake_quant_enabled, self.observer_enabled, self.is_initialized_params,
                   self.quant_min, self.quant_max,
                   self.dtype, self.qscheme, self.ch_axis,
                   self.scale if not self.is_per_channel else 'List',
                   self.zero_point if not self.is_per_channel else 'List')

