import numpy as np
from .BaseOP import FlatbufferBaseOP
from .utils import generate_fake_tensor_info
from mediatools.utils.registry import register_mnn_convert_classes


@register_mnn_convert_classes('AddQuantizePertensorInt')
class BinaryFlatBuffer(FlatbufferBaseOP):
    # main_type
    main_type = "BinaryOp"
    # main commn
    main_optype = 0
    # type
    OPtype = "BinaryOp"

    def __init__(self, node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe,
                        **kwargs) -> None:
        super(BinaryFlatBuffer, self).__init__(**kwargs)
        self.cur_tensor_index = len(extraTensorDescribe)
        self.name = node.output[0]
        self.inputindexes = []
        self.init_params(node, out2node, name2data, extraTensorDescribe)

    def init_params(self, node, out2node, name2data, extraTensorDescribe):
        in_0_n, in_1_n, scale1_n, shift1_n, scale2_n, shift2_n, bias_n, shift_reduce_n = node.input[:]
        self.inputindexes.extend([extraTensorDescribe[in_0_n]['index'], extraTensorDescribe[in_1_n]['index']])
        self.outputOffset = name2data[bias_n].astype(np.int64)[0]
        self.input1Multiplier = name2data[scale1_n].astype(np.int64)[0]
        self.input2Multiplier = name2data[scale2_n].astype(np.int64)[0]
        shift_reduce = name2data[shift_reduce_n].astype(np.int64)[0]
        self.mRightShift1 = name2data[shift1_n].astype(np.int64)[0] - shift_reduce
        self.mRightShift2 = name2data[shift2_n].astype(np.int64)[0] - shift_reduce
        self.outputindexes.append(self.cur_tensor_index)

        qtensor_info = generate_fake_tensor_info(node)
        self.outputActivationMin = qtensor_info['quantInfo']['min']
        self.outputActivationMax = qtensor_info['quantInfo']['max']
        act_quant_bit = int(np.log2(self.outputActivationMax - self.outputActivationMin + 1))
        if self.outputActivationMin < 0:
            self.dtype = "DT_INT8" if act_quant_bit > 4 else "DT_INT4"
        else:
            self.dtype = "DT_UINT8" if act_quant_bit > 4 else "DT_UINT4"

        qtensor_info.update(
            {
                "index": self.cur_tensor_index,
                "name": self.name
            }
        )
        self.qtensor_info.append(qtensor_info)

    def get_main_info(self):
        self.main = {
            "opType": self.main_optype,
            "T": self.dtype,
            "binary": {
                "input1Offset": 0, # fake number
                "input2Offset": 0, # fake number
                "outputOffset": self.outputOffset,
                "input1Multiplier": self.input1Multiplier,
                "input2Multiplier": self.input2Multiplier,
                "outputMultiplier": 0, # fake number
                "outputActivationMin": self.outputActivationMin,
                "outputActivationMax": self.outputActivationMax,
                "mLeftShiftResult1": 0, # fake number
                "mLeftShiftResult2": 0, # fake number
                "mRightShift1": self.mRightShift1,
                "mRightShift2": self.mRightShift2,
                "mLeftShiftOut": 0, # fake number
                "mRightShiftOut": 0, # fake number
                "type": 0 
            }
        }