import os
import math
import numpy as np
import six
import megengine._internal as mgb
from enum import Enum

from py_proto import mace_pb2
from transform import base_converter
from transform.base_converter import PoolingType
from transform.base_converter import ActivationType
from transform.base_converter import EltwiseType
from transform.base_converter import FrameworkType
from transform.base_converter import ReduceType
from transform.base_converter import DataFormat
from transform.base_converter import MaceOp
from transform.base_converter import MaceKeyword
from transform.base_converter import ConverterUtil
from transform.base_converter import RoundMode
from utils.util import mace_check

mge_kernel_h_str = "window_h"
mge_kernel_w_str = "window_w"
mge_stride_h_str = "stride_h"
mge_stride_w_str = "stride_w"
mge_pad_h_str = "pad_h"
mge_pad_w_str = "pad_w"
mge_dilate_h_str = "dilate_h"
mge_dilate_w_str = "dilate_w"

MGESupportedOps = [
    "AxisAddRemove",
    "BatchNormForward",
    "Concat",
    "ConvolutionForward",
    "ConvolutionBackwardData",
    "Dimshuffle",
    "Elemwise",
    "GetVarShape",
    "Host2DeviceCopy",
    "Identity",
    "MarkNoBroadcastElemwise",
    "MatrixMul",
    "PoolingForward",
    "Reduce",
    "Reshape",
    "SharedDeviceTensor",
    "Subtensor",
]

MGEOpType = Enum("MGEOpType", [(op, op) for op in MGESupportedOps], type=str)


def get_symvar_value(mge_symvar):
    if mge_symvar.inferred_value is not None:
        val = mge_symvar.inferred_value
    else:
        cg = mge_symvar.owner_graph
        func = cg.compile_outonly(mge_symvar)
        val = func()

    return val


def is_consumer_group_conv(mge_symvar, var2oprs, map_oprs):
    consumer_ids = var2oprs[mge_symvar.id]
    n_consumers = len(consumer_ids)
    for consumer_id in consumer_ids:
        consumer_op = map_oprs[consumer_id[0]]
        if (
mgb.cgtools.get_opr_type(consumer_op)