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)
                in ("ConvolutionForward", "ConvolutionBackwardData")
                and consumer_op.params["sparse"] == "GROUP"):
            mace_check(n_consumers == 1,
                       "This tensor should only feed depthwise conv/deconv")
            return True
    return False


class MegengineConverter(base_converter.ConverterInterface):
    """A class for convert megengine dumped model to mace model."""

    compute_format_type = {
        "NCHW": DataFormat.NCHW,
        "NHWC": DataFormat.NHWC,
        "DEFAULT": DataFormat.NCHW,
    }

    reduce_math_type = {
        "SUM": ReduceType.SUM,
        "PROD": ReduceType.PROD,
        "MIN": ReduceType.MIN,
        "MAX": ReduceType.MAX,
    }

    # SQE_DIFF, CLIP, SIGN maybe needed
    eltwise_type = {
        "ADD": EltwiseType.SUM,
        "SUB": EltwiseType.SUB,
        "MUL": EltwiseType.PROD,
        "TRUE_DIV": EltwiseType.DIV,
        "MIN": EltwiseType.MIN,
        "MAX": EltwiseType.MAX,
        "NEGATE": EltwiseType.NEG,
        "ABS": EltwiseType.ABS,
        "POW": EltwiseType.POW,
        "EQ": EltwiseType.EQUAL,
        "FLOOR_DIV": EltwiseType.FLOOR_DIV,
        "EXP": EltwiseType.POW,
    }

    activation_type = {
        "RELU": ActivationType.RELU,
        "TANH": ActivationType.TANH,
        "SIGMOID": ActivationType.SIGMOID,
    }

    def __init__(self, option, src_model_file):
        self._op_converters = {
            MGEOpType.AxisAddRemove.name: self.convert_axisaddrm,
            MGEOpType.BatchNormForward.name: self.convert_batchnorm,
            MGEOpType.Concat.name: self.convert_concat,
            MGEOpType.ConvolutionForward.name: self.convert_conv2d,
            MGEOpType.ConvolutionBackwardData.name: self.convert_deconv2d,
            MGEOpType.Dimshuffle.name: self.convert_dimshuffle,
            MGEOpType.Elemwise.name: self.convert_elemwise,
            MGEOpType.GetVarShape.name: self.convert_shape,
            MGEOpType.Host2DeviceCopy.name: self.convert_nop,
            MGEOpType.Identity.name: self.convert_identity,
            MGEOpType.MarkNoBroadcastElemwise.name: self.convert_identity,
            MGEOpType.MatrixMul.name: self.convert_matmul,
            MGEOpType.PoolingForward.name: self.convert_pooling,
            MGEOpType.Reduce.name: self.convert_reduce,
            MGEOpType.Reshape.name: self.convert_reshape,
            MGEOpType.SharedDeviceTensor.name: self.convert_nop,
            MGEOpType.Subtensor.name: self.convert_subtensor,
        }
        self._option = option
        self._mace_net_def = mace_pb2.NetDef()
        ConverterUtil.set_filter_format(self._mace_net_def, DataFormat.OIHW)
        ConverterUtil.add_data_format_arg(self._mace_net_def, DataFormat.NCHW)

        cg, _, outputs = mgb.load_comp_graph_from_file(src_model_file)
        map_oprs, _, var2oprs, *_ = mgb.cgtools.graph_traversal(outputs)
        # prune second input of reshape
        # because it introduces several ops, may increase the overhead
        operators = 
mgb.cgtools.get_oprs_seq(outputs, prune_reshape=True)