import torch 


class QuantizeScheme(object):
    """Describe quantization scheme.
    """
    def __init__(self, symmetry=True, per_channel=False, pot_scale=False, 
                 bit=8, default_bit=None, reduce_range=False, half_range=False,
                 **kwargs):
        self.symmetry = symmetry
        self.per_channel = per_channel
        self.pot_scale = pot_scale
        self.bit = bit
        self.default_bit = default_bit if default_bit is not None else bit
        self.reduce_range = reduce_range
        self.half_range = half_range
        if self.per_channel:
            self.torch_qscheme = torch.per_channel_symmetric if self.symmetry else torch.per_channel_affine
        else:
            self.torch_qscheme = torch.per_tensor_symmetric if self.symmetry else torch.per_tensor_affine
        self.kwargs = kwargs

    def to_observer_params(self):
        naive_para = {
                'dtype': torch.qint8 if not self.half_range else torch.quint8,
                'pot_scale': self.pot_scale,
                'qscheme': self.torch_qscheme,
                'half_range': self.half_range,
                'reduce_range': self.reduce_range,
                'bit': self.bit,
                'ch_axis': 0 if self.per_channel else -1
                # 'default_bit': self.default_bit,
            }
        naive_para.update(self.kwargs)
        return naive_para

    def get_args(self):
        return {
            'symmetry': self.symmetry,
            'per_channel': self.per_channel,
            'pot_scale': self.pot_scale,
            'bit': self.bit,
            # 'default_bit': self.default_bit,
            'reduce_range': self.reduce_range,
            'half_range': self.half_range,
            'ch_axis': 0 if self.per_channel else -1
        }

    def __str__(self):
        return "Symmetric: {} / Bitwidth: {} / DefaultBit: {} / Per channel: {} / Pot scale: {} / Reduce range: {}".format(self.symmetry, 
                                                                                       self.bit,
                                                                                       self.default_bit,
                                                                                       self.per_channel,
                                                                                       self.pot_scale,
                                                                                       self.reduce_range)
