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


@register_mnn_convert_classes('QuantizePertensorInt')
@register_mnn_convert_classes('QuantizePerchannelInt')
class InputFlatBuffer(FlatbufferBaseOP):
    # main_type
    main_type = "Input"
    # main commn
    dtype = "DT_UINT8"
    dformat = "NC4HW4"
    # type
    OPtype = "Input"
    def __init__(self, node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe,
                        **kwargs) -> None:
        super(InputFlatBuffer, self).__init__(**kwargs)
        self.cur_tensor_index = len(extraTensorDescribe)
        self.init_params(node, out2node)

    def init_params(self, node, out2node):
        in_node = out2node[node.input[0]]
        self.name = node.output[0]
        self.dims = [x.dim_value for x in in_node.type.tensor_type.shape.dim]
        self.outputindexes.append(self.cur_tensor_index)
        qtensor_info = generate_fake_tensor_info(node)
        qtensor_info.update(
            {
                "index": self.cur_tensor_index,
                "name": self.name
            }
        )
        self.qtensor_info.append(qtensor_info)

    def get_main_info(self):
        self.main = {
            "dims": self.dims,
            "dtype": self.dtype,
            "dformat": self.dformat
        }
    
