class QuantizedConv(nn.HybridBlock):
    def __init__(self, 
                 kernel_size, 
                 strides, 
                 num_filter, 
                 kernel_shape,
                 activation=None,
                 dilation=(1, 1),
                 padding=(0, 0),
                 groups=1,
                 use_bias=True,
                 layout='NCHW',
                 adj=None,
                 in_channels=0,
                 op_name='Convolution',
                 weight_initializer=None,
                 weight_min_initializer=None,
                 weight_max_initializer=None,
                 bias_initializer='zeros',
                 weight_shape=None,
                 **kwargs):
        super(QuantizedConv, self).__init__(**kwargs)
        self._kwargs = {
                'kernel': kernel_size, 'stride': strides, 'dilate': dilation,
                'pad': padding, 'num_filter': num_filter, 'num_group': groups,
                'no_bias': not use_bias, 'layout': layout}
        if adj is not None:
            self._kwargs['adj'] = adj
        dshape = [0]*(len(kernel_size) + 2)
        dshape[layout.find('N')] = 1
        dshape[layout.find('C')] = in_channels
        if weight_shape is None:
            self.wshapes = _infer_weight_shape(op_name, dshape, self._kwargs)
        else:
            self.wshapes = weight_shape
        
        self.weight = self.params.get('weight', shape=kernel_shape,
                                          init=weight_initializer,
                                          allow_deferred_init=True)
        self.weight_min = self.params.get('weight_min', shape=(1,),
                                         init=weight_min_initializer,
                                         allow_deferred_init=True)
        self.weight_max = self.params.get('weight_max', shape=(1,),
                                         init=weight_max_initializer,
                                         allow_deferred_init=True)
        
        if use_bias:
            self.bias = self.params.get('bias', shape=self.wshapes[2],
                                            init=bias_initializer,
                                            allow_deferred_init=True)
            self.bias_min = self.params.get('bias_min', shape=(1,),
                                            init=bias_initializer,
                                            allow_deferred_init=True)
            self.bias_max = self.params.get('bias_max', shape=(1,),
                                            init=bias_initializer,
                                            allow_deferred_init=True)
        else:
            self.bias = None
            self.bias_min = None
            self.bias_max = None

        self.kernel_size = kernel_size
        self.strides = strides
        self.num_filter = num_filter
        # for string reprentation
        self.layout = layout
        
        if activation is not None:
            self.act = Activation(activation, prefix=activation+'_')
        else:
            self.act = None

        
    def hybrid_forward(self, F, x, weight, weight_min, weight_max, bias=None, 
                       bias_min=None, bias_max=None):
        if not isinstance(weight, mxnet.ndarray.ndarray.NDArray): 
            print('F={}'.format(F))
            weight_temp = weight.astype(dtype='int8')
            print(weight)
            print(weight_temp)
            print(F.contrib.requantize(weight_temp, weight.min(), weight.max()))
        q_inputs, q_inputs_min, q_inputs_max = F.contrib.quantize(
            data=x, 
            min_range=x.min(), 
            max_range=x.max(),
            out_type='uint8')
        q_output, q_output_min, q_output_max = F.contrib.quantized_conv(
                       data=q_inputs, 
                       weight=weight.astype(dtype ='int8'), 
                       bias=bias.astype(dtype='int8'), 
                       min_data=q_inputs_min, 
                       max_data=q_inputs_max, 
                       min_weight=weight_min,
                       max_weight=weight_max,
                       min_bias=bias_min,
                       max_bias=bias_max, 
                       kernel=self.kernel_size,
                       stride=self.strides,
                       num_filter=self.num_filter, 
                       layout=self.layout,
                       cudnn_off=True, 
                       cudnn_tune='off',
                       name='fwd')
        q_8_out, q_8_out_min, q_8_out_max = F.contrib.requantize(
            data=q_output.astype('int32'), 
            min_range=q_output_min, 
            max_range=q_output_max,
            name='fwd', 
        )
        act = F.contrib.dequantize(
            data=q_8_out, 
            min_range=q_8_out_min, 
            max_range=q_8_out_max,
            name='fwd')
        if self.act is not None:
            act = self.act(act)
        return act
    
    def __repr__(self):
        s = '{name}({mapping}, kernel_size={kernel}, stride={stride}'
        len_kernel_size = len(self._kwargs['kernel'])
        if self._kwargs['pad'] != (0,) * len_kernel_size:
            s += ', padding={pad}'
        if self._kwargs['dilate'] != (1,) * len_kernel_size:
            s += ', dilation={dilate}'
        if hasattr(self, 'out_pad') and self.out_pad != (0,) * len_kernel_size:
            s += ', output_padding={out_pad}'.format(out_pad=self.out_pad)
        if self._kwargs['num_group'] != 1:
            s += ', groups={num_group}'
        if self.weight is None:
            s += ', bias=False'
        if self.act:
            s += ', {}'.format(self.act)
        s += ')'
        shape = self.weight.shape
        return s.format(name=self.__class__.__name__,
                        mapping='{0} -> {1}'.format(shape[1] if shape[1] else None, shape[0]),
                        **self._kwargs)