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('Conv2dFixpoint')
class Conv2dFlatbuffer(FlatbufferBaseOP):
    out_ch_axis = 0
    # main_type
    main_type = "Convolution2D"
    # main commn
    padMode = "CAFFE"
    is_relu = False
    is_relu6 = False
    hasOutputShape = False
    # main quanParameter
    type = 4
    useInt32 = False
    readType = 0
    has_scaleInt = False
    # main symmetricQuan
    method = "DEFAULT"
    # type
    OPtype = "Convolution"

    def __init__(self, node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe, 
                        **kwargs) -> None:
        super(Conv2dFlatbuffer, self).__init__(**kwargs)
        self.cur_tensor_index = len(extraTensorDescribe)
        self.inputindexes = []
        self.init_op_params(node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe)

    def init_op_params(self, node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe):
        self.init_conv_params(node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe)
        self.init_quant_params(node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe)

    def init_quant_params(self, node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe):
        assert len(inp2node[node.output[0]]) == 1, "fixpoint conv must be followed with only one MixedQuantize Node"
        next_node = inp2node[node.output[0]][0][0]
        self.name = next_node.output[0]
        self.qt_info = generate_fake_tensor_info(next_node)
        scale_n, shift_n, bias_n, clip_value_n = next_node.input[1:5]
        self.mquantizeMultiplier_all = name2data[scale_n].astype(np.int64)
        if len(self.mquantizeMultiplier_all.tolist()) == 1:
            self.mquantizeMultiplier_all = self.mquantizeMultiplier_all.repeat([self.outputCount])
        self.mRightShift_all = name2data[shift_n].astype(np.int64)
        if self.mRightShift_all.shape[0] == 1:
            self.mRightShift_all = self.mRightShift_all.repeat(self.outputCount, axis=0)
        self.outputZeroPoint = name2data[clip_value_n].astype(np.int64)[0]
        self.fake_float_scale_all = np.zeros_like(self.mquantizeMultiplier_all, dtype=np.float32)
        self.clamp_qmin = self.qt_info['quantInfo']['min']
        self.clamp_qmax = self.qt_info['quantInfo']['max']

    def init_conv_params(self, node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe):
        input_name, self.weight_name, self.bias_name = node.input
        assert input_name in extraTensorDescribe, "tensor name not found!!!"
        self.inputindexes.append(extraTensorDescribe[input_name]['index'])
        self.weight_init = named_initializer[self.weight_name]
        self.bias_init = named_initializer[self.bias_name]
        self.outputCount = self.weight_init.dims[0]
        self.inputCount = self.weight_init.dims[1]
        for _, attr in enumerate(node.attribute):
            if attr.name in ['kernel_shapes', 'strides', 'pads', 'dilations']:
                setattr(self, attr.name, attr.ints)
            elif attr.name in ['group', 'is_relu', 'padding_value', 'weight_bitwidth']:
                setattr(self, attr.name, attr.i)
            else:
                raise NameError('Unknown attribute name: %s'%(attr.name))
        
        if self.inputCount == 1 and self.group != 1:
            self.OPtype = "ConvolutionDepthwise"
            self.inputCount = self.group
        self.fixpoint_bias_all = numpy_helper.to_array(self.bias_init).astype(np.int64)
        self.fixpoint_weight_all = numpy_helper.to_array(self.weight_init).astype(np.int64)
        if self.out_ch_axis == 1:
            # for deconv transpose out ch
            self.fixpoint_weight = self.fixpoint_weight.transpose((1, 0, 2, 3))
        self.fake_float_bias_all = np.zeros_like(self.fixpoint_bias_all)

    def get_main_info(self):
        self.common_info = {
            'padX': self.pads[0],
            'padY': self.pads[1],
            'kernelX': self.kernel_shapes[0],
            'kernelY': self.kernel_shapes[1],
            'strideX': self.strides[0],
            'strideY': self.strides[1],
            'dilateX': self.dilations[0],
            'dilateY': self.dilations[1],
            'padMode': self.padMode,
            'group': self.group,
            'outputCount': self.outputCount,
            'inputCount': self.inputCount,
            'relu': self.is_relu,
            'relu6': self.is_relu6,
            'pads': 2*list(self.pads),
            'hasOutputShape': self.hasOutputShape
        }
        self.main = {
            "common": self.common_info,
            "bias": self.fake_float_bias.tolist(),
            "quanParameter": {
                "buffer": self.fixpoint_weight.reshape(-1).tolist(),
                "alpha": self.fake_float_scale.tolist(),
                "type": self.type,
                "useInt32": self.useInt32,
                "quantScale": 1.0, # fake number
                "scaleIn": 1.0, # fake number
                "scaleOut": 1.0, # fake number
                "aMax": self.outputCount,
                "aMin": 0,
                "readType": self.readType,
                "has_scaleInt": self.has_scaleInt
            },
            "symmetricQuan": {
                "method": self.method,
                "nbits": self.weight_bitwidth,
                "zeroPoint": self.padding_value,
                "outputZeroPoint": self.outputZeroPoint,
                "clampMin": self.clamp_qmin,
                "clampMax": self.clamp_qmax
            },
            "convolution2d_superacme": {
                "mquantizeMultiplier": self.mquantizeMultiplier.tolist() ,
                "mRightShift": self.mRightShift.tolist(),
                "mintbias": self.fixpoint_bias.tolist(),
                "mmode": 4,
                "weightType": "DT_INT8" if self.weight_bitwidth > 4 else "DT_INT4",
            },
            "weightname": self.weight_name
        }

    def dump2json(self):
        if self.group == 1 or self.OPtype == "ConvolutionDepthwise":
            self.outputindexes.append(self.cur_tensor_index)
            self.qt_info.update(
                {
                    "index": self.cur_tensor_index,
                    "name": self.name
                }
            )
            self.qtensor_info.append(self.qt_info)
            self.fixpoint_bias = self.fixpoint_bias_all
            self.fixpoint_weight = self.fixpoint_weight_all
            self.mquantizeMultiplier = self.mquantizeMultiplier_all
            self.mRightShift = self.mRightShift_all
            self.fake_float_scale = self.fake_float_scale_all
            self.fake_float_bias = self.fake_float_bias_all
            return super(Conv2dFlatbuffer, self).dump2json()
        else:
            # for group conv we insert slice and concat op
            return self.dump2json_groupconv()

    def dump2json_groupconv(self):
        op_infos = []
        qtensor_info = []
        # add slice first
        assert self.group == 2, "not supported for group > 2"
        op_infos.append(
            {
            "inputIndexes": self.inputindexes,
            "main_type": "Slice",
            "main": {
                "axis": 1,
                "slicePoints": [
                self.inputCount
                ]
            },
            "name": "%s__slice"%(self.name),
            "outputIndexes": [
                self.cur_tensor_index,
                self.cur_tensor_index + 1
            ],
            "type": "Slice"
            }
        )
        for i in range(self.group):
            cur_qtensor_info = self.qt_info.copy()
            cur_qtensor_info.update({
                'index': self.cur_tensor_index + i,
                'name': '%s__input__%s'%(self.name, i)
            })
            qtensor_info.append(cur_qtensor_info)
        # group conv info
        ori_name = self.name
        group = self.group
        group_step = self.outputCount // group
        for i in range(group):
            self.fixpoint_bias = self.fixpoint_bias_all[i*group_step:(i+1)*group_step]
            self.fixpoint_weight = self.fixpoint_weight_all[i*group_step:(i+1)*group_step, :, :, :]
            self.mquantizeMultiplier = self.mquantizeMultiplier_all[i*group_step:(i+1)*group_step]
            self.mRightShift = self.mRightShift_all[i*group_step:(i+1)*group_step]
            self.fake_float_scale = self.fake_float_scale_all[i*group_step:(i+1)*group_step]
            self.fake_float_bias = self.fake_float_bias_all[i*group_step:(i+1)*group_step]
            self.inputindexes = [self.cur_tensor_index + i]
            self.outputindexes = [self.cur_tensor_index + len(qtensor_info)]
            self.outputCount = group_step
            self.group = 1
            self.name = "%s__group__%s"%(ori_name, i)
            self.get_main_info()
            op_infos.append(self.get_op_info())

            cur_qtensor_info = self.qt_info.copy()
            cur_qtensor_info.update({
                'index': self.cur_tensor_index + 1,
                'name': '%s__output__%s'%(self.name, i)
            })
            qtensor_info.append(cur_qtensor_info)

        # add concat op
        op_infos.append(
            {
            "inputIndexes": [
                op_infos[-2]['outputIndexes'][0],
                op_infos[-1]['outputIndexes'][0]
            ],
            "main_type": "Axis",
            "main": {
                "axis": 1
            },
            "name": ori_name,
            "outputIndexes": [
                self.cur_tensor_index + len(qtensor_info)
            ],
            "type": "Concat"
            }
        )
        cur_qtensor_info = self.qt_info.copy()
        cur_qtensor_info.update({
            'index': self.cur_tensor_index + len(qtensor_info),
            'name': ori_name
        })
        qtensor_info.append(cur_qtensor_info)
        return op_infos, qtensor_info

