# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, import-self
"""MXNet symbol frontend."""
from __future__ import absolute_import as _abs
import json
import tvm
from .. import symbol as _sym
from .common import get_nnvm_op, required_attr, parse_tshape, parse_bool_str

__all__ = ['from_mxnet']

def _rename(new_name):
    def impl(inputs, attrs):
        return get_nnvm_op(new_name)(*inputs, **attrs)
    return impl

def _pooling(inputs, attrs):
    kernel = parse_tshape(required_attr(attrs, 'kernel', 'pooling'))
    if len(kernel) != 2:
        raise tvm.error.OpAttributeUnimplemented(
            'Non-2D kernels are not supported for Pool2D.')
    global_pool = 'global' if parse_bool_str(attrs, 'global_pool') else ''
    pool_type = required_attr(attrs, 'pool_type', 'pooling')
    if pool_type not in ['avg', 'max']:
        raise tvm.error.OpNotImplemented(
            'Only max and average pooling are supported in frontend MXNet.')
    op_name, new_attrs = '_'.join([global_pool, pool_type, 'pool2d']).strip('_'), {}
    # new_attrs['layout'] = 'NCHW'
    if not global_pool:
        new_attrs['pool_size'] = kernel
        new_attrs['strides'] = attrs.get('stride', (1, 1))
        new_attrs['padding'] = attrs.get('pad', (0, 0))
        new_attrs['ceil_mode'] = (attrs.get('pooling_convention', 'valid') == 'full')
        if pool_type == 'avg':
            new_attrs['count_include_pad'] = attrs.get('count_include_pad', True)
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _batch_norm(inputs, attrs):
    if parse_bool_str(attrs, 'output_mean_var'):
        raise tvm.error.OpAttributeUnimplemented(
            'Attribute "output_mean_var" is not supported in operator batch_norm.')
    # if parse_bool_str(attrs, 'fix_gamma'):
    #     _warn_not_used('fix_gamma', 'batch_norm')
    if parse_bool_str(attrs, 'use_global_stats'):
        from warnings import warn
        warn(
            'Attribute "use_global_stats" is ignored in operator batch_norm.')
    # if parse_bool_str(attrs, 'momentum'):
    #     _warn_not_used('momentum', 'batch_norm')
    op_name, new_attrs = 'batch_norm', {}
    new_attrs['axis'] = attrs.get('axis', 1)
    new_attrs['epsilon'] = attrs.get('eps', 0.001)
    new_attrs['center'] = True
    new_attrs['scale'] = not parse_bool_str(attrs, 'fix_gamma', default="False")
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _concat(inputs, attrs):
    op_name = 'concatenate'
    new_attrs = {'axis': attrs.get('dim', 1)}
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _conv2d(inputs, attrs):
    kernel = parse_tshape(required_attr(attrs, 'kernel', 'conv2d'))
    if len(kernel) != 2:
        raise tvm.error.OpAttributeUnimplemented(
            'Non-2D kernels are not supported for operator Conv2D.')
    layout = attrs.get('layout', 'NCHW')
    if layout not in ['NCHW', 'NHWC']:
        raise tvm.error.OpAttributeUnimplemented(
            'Layout {} is not supported in operator Conv2D.'.format(layout))
    if 'kernel_layout' in attrs:
        kernel_layout = attrs['kernel_layout']
    else:
        kernel_layout = 'HWIO' if layout == 'NHWC' else 'OIHW'
    op_name, new_attrs = 'conv2d', {}
    new_attrs['channels'] = required_attr(attrs, 'num_filter', 'conv2d')
    new_attrs['kernel_size'] = kernel
    new_attrs['strides'] = attrs.get('stride', (1, 1))
    new_attrs['padding'] = attrs.get('pad', (0, 0))
    new_attrs['dilation'] = attrs.get('dilate', (1, 1))
    new_attrs['groups'] = attrs.get('num_group', 1)
    new_attrs['layout'] = layout
    new_attrs['kernel_layout'] = kernel_layout
    new_attrs['use_bias'] = attrs.get('no_bias', 'False').strip() == 'False'
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _conv2d_transpose(inputs, attrs):
    if 'target_shape' in attrs:
        raise tvm.error.OpAttributeUnimplemented(
            'Attribute "target_shape" is not supported in operator Conv2D-transpose.')
    kernel = parse_tshape(required_attr(attrs, 'kernel', 'conv2d_transpose'))
    if len(kernel) != 2:
        raise tvm.error.OpAttributeInvalid(
            'Non-2D kernels are not supported in Conv2D-transpose.')
    layout = attrs.get('layout', 'NCHW')
    if layout not in ['NCHW', 'NHWC']:
        raise tvm.error.OpAttributeUnimplemented(
            'Layout {} is not supported in operator Conv2D-transpose.')
    if 'kernel_layout' in attrs:
        kernel_layout = attrs['kernel_layout']
    else:
        kernel_layout = 'HWIO' if layout == 'NHWC' else 'OIHW'
    op_name, new_attrs = 'conv2d_transpose', {}
    new_attrs['channels'] = required_attr(attrs, 'num_filter', 'conv2d_transpose')
    new_attrs['kernel_size'] = kernel
    new_attrs['strides'] = attrs.get('stride', (1, 1))
    new_attrs['output_padding'] = attrs.get('adj', (0, 0))
    new_attrs['padding'] = attrs.get('pad', (0, 0))
    new_attrs['dilation'] = attrs.get('dilate', (1, 1))
    new_attrs['groups'] = attrs.get('num_group', 1)
    new_attrs['layout'] = layout
    new_attrs['kernel_layout'] = kernel_layout
    new_attrs['use_bias'] = not parse_bool_str(attrs, 'no_bias')
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _dense(inputs, attrs):
    import mxnet as mx
    op_name, new_attrs = 'dense', {}
    new_attrs['units'] = required_attr(attrs, 'num_hidden', 'dense')
    new_attrs['use_bias'] = not parse_bool_str(attrs, 'no_bias')
    try:
        _ = mx.sym.FullyConnected(mx.sym.var('x'), num_hidden=1, flatten=True)
        has_flatten = True
    except mx.base.MXNetError:
        # no flatten attribute in old mxnet
        has_flatten = False
    use_flatten = parse_bool_str(attrs, 'flatten', 'True')
    if has_flatten and use_flatten:
        inputs[0] = _sym.flatten(inputs[0])
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _dropout(inputs, attrs):
    op_name, new_attrs = 'dropout', {}
    new_attrs['rate'] = attrs.get('p', 0.5)
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _leaky_relu(inputs, attrs):
    act_type = required_attr(attrs, 'act_type', 'leaky_relu')
    if act_type in ['leaky', 'prelu']:
        op_name, new_attrs = act_type, {}
        if act_type == 'leaky':
            new_attrs['alpha'] = attrs.get('slope', 0.25)
        sym = get_nnvm_op(op_name)(*inputs, **new_attrs)
    elif act_type == 'elu':
        slope = attrs.get('slope', 0.25)
        sym = -slope * _sym.relu(1 - _sym.exp(*inputs)) + _sym.relu(*inputs)
    elif act_type == 'rrelu':
        lower_bound = float(required_attr(attrs, 'lower_bound', 'leaky_relu'))
        upper_bound = float(required_attr(attrs, 'upper_bound', 'leaky_relu'))
        slope = (lower_bound + upper_bound) / 2.0
        op_name, new_attrs = 'leaky_relu', {'alpha': str(slope)}
        sym = get_nnvm_op(op_name)(*inputs, **new_attrs)
    else:
        raise tvm.error.OpNotImplemented(
            'Operator {} is not supported in frontend MXNet.'.format(act_type))
    return sym

def _activations(inputs, attrs):
    act_type = required_attr(attrs, 'act_type', 'activations')
    if act_type in ['relu', 'sigmoid', 'tanh']:
        op_name, new_attrs = act_type, {}
        sym = get_nnvm_op(op_name)(*inputs, **new_attrs)
    elif act_type == 'softrelu':
        sym = _sym.log((1 + _sym.exp(*inputs)))
    else:
        raise tvm.error.OpNotImplemented(
            'Operator {} is not supported in frontend MXNet.'.format(act_type))
    return sym

def _reshape(inputs, attrs):
    if parse_bool_str(attrs, 'reverse'):
        raise tvm.error.OpAttributeUnimplemented(
            'Attribute "reverse" is not supported in operator Reshape.')
    op_name, new_attrs = 'reshape', {}
    new_attrs['shape'] = required_attr(attrs, 'shape', 'reshape')
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _slice(inputs, attrs):
    begin = attrs.get('begin', None)
    end = attrs.get('end', None)
    stride = attrs.get('step', None)
    if begin is None or end is None:
        raise RuntimeError('begin and end are required params')
    if 'None' in begin or 'None' in end:
        raise RuntimeError('None in begin or end not supported yet...')
    new_attrs = {'begin': begin, 'end': end}
    if stride is not None:
        new_attrs['stride'] = stride
    return get_nnvm_op('strided_slice')(inputs[0], **new_attrs)

def _split(inputs, attrs):
    op_name, new_attrs = 'split', {}
    axis = attrs.get('axis', 1)
    new_attrs['indices_or_sections'] = required_attr(attrs, 'num_outputs', 'split')
    new_attrs['axis'] = axis
    outputs = get_nnvm_op(op_name)(*inputs, **new_attrs)
    if parse_bool_str(attrs, 'squeeze_axis'):
        squeeze_attrs = {'axis': axis}
        outputs = _sym.Group([get_nnvm_op('squeeze')(o, **squeeze_attrs) for o in outputs])
    return outputs

def _softmax_activation(inputs, attrs):
    op_name, new_attrs = 'softmax', {}
    mode = attrs.get('mode', 'instance')
    new_attrs['axis'] = 0 if mode == 'instance' else 1
    return get_nnvm_op(op_name)(inputs[0], **new_attrs)

def _softmax_output(inputs, attrs):
    op_name, new_attrs = 'softmax', {}
    if parse_bool_str(attrs, 'multi_output'):
        new_attrs['axis'] = 1
    return get_nnvm_op(op_name)(inputs[0], **new_attrs)

def _upsampling(inputs, attrs):
    scale = attrs.get('scale')
    new_attrs = {'scale':int(scale)}
    return get_nnvm_op('upsampling')(inputs[0], **new_attrs)

def _clip(inputs, attrs):
    op_name, new_attrs = "clip", {}
    new_attrs['a_min'] = required_attr(attrs, 'a_min', 'clip')
    new_attrs['a_max'] = required_attr(attrs, 'a_max', 'clip')
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _contrib_multibox_detection(inputs, attrs):
    clip = parse_bool_str(attrs, 'clip', default='True')
    threshold = attrs.get('threshold') or 0.01
    nms_threshold = attrs.get('nms_threshold') or 0.5
    force_suppress = parse_bool_str(attrs, 'force_suppress', default='False')
    variances = tuple([float(x.strip()) for x in attrs.get('variances').strip('()').split(',')]) \
        if attrs.get('variances') is not None else (0.1, 0.1, 0.2, 0.2)
    nms_topk = attrs.get('nms_topk') or -1
    new_attrs0 = {'clip': clip, 'threshold': float(threshold), 'variances': variances}
    new_attrs1 = {'return_indices': False, 'iou_threshold': float(nms_threshold),
                  'force_suppress': force_suppress, 'top_k': int(nms_topk)}
    data, valid_count = get_nnvm_op('multibox_transform_loc')(inputs[0], inputs[1],
                                                              inputs[2], **new_attrs0)
    return get_nnvm_op('non_max_suppression')(data, valid_count, **new_attrs1)

def _elemwise_sum(inputs, _):
    new_attrs = {'num_args':len(inputs)}
    return get_nnvm_op('elemwise_sum')(*inputs, **new_attrs)

def _crop_like(inputs, attrs):
    new_attrs = {}
    offsets = \
        tuple([float(x.strip()) for x in attrs.get('offsets').strip('()').split(',')]) \
            if attrs.get('offsets') is not None else (0, 0)
    if offsets != (0, 0):
        raise tvm.error.OpAttributeInvalid(
            'crop_like offsets must equal (0,0).')
    center_crop = parse_bool_str(attrs, 'center_crop', default="False")
    if center_crop:
        raise tvm.error.OpAttributeUnimplemented(
            'Center crop is not supported in operator crop_like.')
    if len(inputs) < 2:
        raise RuntimeError("Only support crop_like pattern.")
    new_attrs["axis"] = [2, 3]
    return get_nnvm_op('slice_like')(inputs[0], inputs[1], **new_attrs)


def _expand_dims(inputs, attrs):
    op_name, new_attrs = 'expand_dims', {}
    new_attrs['axis'] = required_attr(attrs, 'axis', 'expand_dims')
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _lrn(inputs, attrs):
    op_name, new_attrs = 'lrn', {}
    new_attrs['alpha'] = attrs.get('alpha', 0.0001)
    new_attrs['beta'] = attrs.get('beta', 0.75)
    new_attrs['bias'] = attrs.get('knorm', 2)
    # NCHW format and normalization along channel axis
    new_attrs['axis'] = 1
    new_attrs['size'] = required_attr(attrs, 'nsize', 'lrn')
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _minimum(inputs, attrs):
    return get_nnvm_op('broadcast_min')(*inputs, **attrs)

def _maximum(inputs, attrs):
    return get_nnvm_op('broadcast_max')(*inputs, **attrs)

def _ones(_, attrs):
    op_name = 'ones'
    return get_nnvm_op(op_name)(**attrs)

def _zeros(_, attrs):
    op_name = 'zeros'
    return get_nnvm_op(op_name)(**attrs)

def _argmax(inputs, attrs):
    op_name, new_attrs = 'argmax', {}
    new_attrs['dtype'] = 'float32'
    new_attrs['axis'] = attrs.get('axis', 0)
    new_attrs['keepdims'] = parse_bool_str(attrs, 'keepdims', default="False")
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

def _argmin(inputs, attrs):
    op_name, new_attrs = 'argmin', {}
    new_attrs['dtype'] = 'float32'
    new_attrs['axis'] = attrs.get('axis', 0)
    new_attrs['keepdims'] = parse_bool_str(attrs, 'keepdims', default="False")
    return get_nnvm_op(op_name)(*inputs, **new_attrs)

_identity_list = ['__add_scalar__', '__add_symbol__', '__div_scalar__',
                  '__div_symbol__', '__mul_scalar__', '__mul_symbol__',
                  '__pow_scalar__', '__rdiv_scalar__', '__rpow_scalar__',
                  '__rsub_scalar__', '__sub_scalar__', '__sub_symbol__',
                  'broadcast_add', 'broadcast_div', 'broadcast_mul',
                  'broadcast_sub', 'broadcast_to', 'cast', 'elemwise_add',
                  'elemwise_div', 'elemwise_mul', 'elemwise_sub', 'exp',
                  'flatten', 'log', 'log_softmax', 'max', 'min', 'negative',
                  'ones_like', 'relu', 'sigmoid', 'slice_like', 'softmax',
                  'sum', 'tanh', 'transpose', 'zeros_like', 'gather_nd',
                  'reshape_like', 'where']

_convert_map = {
    '_copy'         : _rename('copy'),
    '_div_scalar'   : _rename('__div_scalar__'),
    '_minus_scalar' : _rename('__sub_scalar__'),
    '_mul_scalar'   : _rename('__mul_scalar__'),
    '_plus_scalar'  : _rename('__add_scalar__'),
    '_rdiv_scalar'  : _rename('__rdiv_scalar__'),
    '_rminus_scalar': _rename('__rsub_scalar__'),
    '_contrib_MultiBoxPrior' : _rename('multibox_prior'),
    '_contrib_MultiBoxDetection' : _contrib_multibox_detection,
    '_minimum'      : _minimum,
    '_maximum'      : _maximum,
    '_ones'         : _ones,
    '_zeros'        : _zeros,
    'argmax'        : _argmax,
    'argmin'        : _argmin,
    'Activation'    : _activations,
    'BatchNorm'     : _batch_norm,
    'BatchNorm_v1'  : _batch_norm,
    'Cast'          : _rename('cast'),
    'Concat'        : _concat,
    'Convolution'   : _conv2d,
    'Convolution_v1': _conv2d,
    'Crop'          : _crop_like,
    'Deconvolution' : _conv2d_transpose,
    'Dropout'       : _dropout,
    'Flatten'       : _rename('flatten'),
    'FullyConnected': _dense,
    'LeakyReLU'     : _leaky_relu,
    'Pooling'       : _pooling,
    'Pooling_v1'    : _pooling,
    'Reshape'       : _reshape,
    'slice'         : _slice,
    'SliceChannel'  : _split,
    'split'         : _split,
    'Softmax'       : _rename('softmax'),
    'SoftmaxActivation' : _softmax_activation,
    'SoftmaxOutput' : _softmax_output,
    'add_n'         : _elemwise_sum,
    'concat'        : _concat,
    'max_axis'      : _rename('max'),
    'min_axis'      : _rename('min'),
    'reshape'       : _reshape,
    'sum_axis'      : _rename('sum'),
    'UpSampling'    : _upsampling,
    'clip'          : _clip,
    'expand_dims'   : _expand_dims,
    'LRN'           : _lrn
}

def _convert_symbol(op_name, inputs, attrs,
                    identity_list=None,
                    convert_map=None):
    """Convert from mxnet op to nnvm op.
    The converter must specify some conversions explicitly to
    support gluon format ops such as conv2d...

    Parameters
    ----------
    op_name : str
        Operator name, such as Convolution, FullyConnected
    inputs : list of nnvm.Symbol
        List of input symbols.
    attrs : dict
        Dict of operator attributes
    identity_list : list
        List of operators that don't require conversion
    convert_map : dict
        Dict of name : callable, where name is the op's name that
        require conversion to nnvm, callable are functions which
        take attrs and return (new_op_name, new_attrs)

    Returns
    -------
    sym : nnvm.Symbol
        Converted nnvm Symbol
    """
    identity_list = identity_list if identity_list else _identity_list
    convert_map = convert_map if convert_map else _convert_map
    if op_name in identity_list:
        op = get_nnvm_op(op_name)
        sym = op(*inputs, **attrs)
    elif op_name in convert_map:
        sym = convert_map[op_name](inputs, attrs)
    else:
        raise tvm.error.OpNotImplemented(
            'Operator {} is not supported in frontend MXNet.'.format(op_name))
    return sym

def _as_list(arr):
    """Force being a list, ignore if already is."""
    if isinstance(arr, list):
        return arr
    return [arr]

def _topo_sort(symbol):
    """Sort all symbols in the mxnet graph in topological order.

    Parameters
    ----------
    symbol : mxnet.sym.Symbol

    Returns:
    -------
    list
        List of mxnet symbol
    """
    queue = []
    symbol_map = {}
    deps = {}
    dep_cnts = {}
    for s in symbol:
        symbol_map[s.attr('name')] = s
        queue.append(s)
    while queue:
        sym = queue.pop(0)
        name = sym.attr('name')
        childs = sym.get_children()
        if childs is None:
            dep_cnts[name] = 0
        else:
            dep_cnts[name] = len({c.attr('name') for c in childs})
            for child in childs:
                child_name = child.attr('name')
                if child_name not in deps:
                    deps[child_name] = set()
                deps[child_name].add(name)
                if child_name not in symbol_map:
                    symbol_map[child_name] = child
                    queue.append(child)
    order = []
    while dep_cnts:
        remove = []
        for name in dep_cnts:
            if dep_cnts[name] == 0:
                order.append(symbol_map[name])
                remove.append(name)
                if name in deps:
                    for other in deps[name]:
                        dep_cnts[other] -= 1
        for name in remove:
            del dep_cnts[name]
    return order

def _from_mxnet_impl(symbol, graph):
    """Convert mxnet symbol to nnvm implementation.
    Reconstruct a nnvm symbol by traversing the mxnet symbol.

    Parameters
    ----------
    symbol : mxnet.sym.Symbol
        Incompatible symbol from mxnet, sharing similar graph structure.
        The op_name and attrs inside are not always compatible.
    graph : dict
        Reusable nodes are stored in graph.

    Returns:
    -------
    nnvm.sym.Symbol
        Converted symbol
    """
    def get_node(sym):
        name = sym.attr('name')
        if name not in graph:
            return None
        output_index = json.loads(sym.tojson())['heads'][0][1]
        return graph[name][output_index]

    assert symbol is not None
    # Traverse all symbols in topological order
    for sym in _topo_sort(symbol):
        name = sym.attr('name')
        attr = sym.list_attr()
        op_name = sym.attr('op_name')
        childs = sym.get_children()
        if childs is not None:
            childs = [get_node(child) for child in childs]
            childs = [x for y in childs for x in _as_list(y)]
            node = _convert_symbol(op_name, childs, attr)
        elif op_name != 'null':
            node = _convert_symbol(op_name, [], attr)
        else:
            node = _sym.Variable(name=name, **attr)
        graph[name] = node
    nodes = []
    for sym in symbol:
        node = get_node(sym)
        assert node is not None
        nodes.append(node)
    if len(nodes) > 1:
        return _sym.Group(nodes)
    return nodes[0]

def from_mxnet(symbol, arg_params=None, aux_params=None):
    """Convert from MXNet's model into compatible NNVM format.

    Parameters
    ----------
    symbol : mxnet.Symbol or mxnet.gluon.HybridBlock
        MXNet symbol

    arg_params : dict of str to mx.NDArray
        The argument parameters in mxnet

    aux_params : dict of str to mx.NDArray
        The auxiliary parameters in mxnet

    Returns
    -------
    sym : nnvm.Symbol
        Compatible nnvm symbol

    params : dict of str to tvm.NDArray
        The parameter dict to be used by nnvm
    """
    try:
        import mxnet as mx
    except ImportError as e:
        raise ImportError('{}. MXNet is required to parse symbols.'.format(e))

    if isinstance(symbol, mx.sym.Symbol):
        sym = _from_mxnet_impl(symbol, {})
        params = {}
        arg_params = arg_params if arg_params else {}
        aux_params = aux_params if aux_params else {}
        for k, v in arg_params.items():
            params[k] = tvm.nd.array(v.asnumpy())
        for k, v in aux_params.items():
            params[k] = tvm.nd.array(v.asnumpy())
    elif isinstance(symbol, mx.gluon.HybridBlock):
        data = mx.sym.Variable('data')
        sym = symbol(data)
        sym = _from_mxnet_impl(sym, {})
        params = {}
        for k, v in symbol.collect_params().items():
            params[k] = tvm.nd.array(v.data().asnumpy())
    elif isinstance(symbol, mx.gluon.Block):
        raise NotImplementedError("Only Hybrid Blocks are supported now.")
    else:
        msg = "mxnet.Symbol or gluon.HybridBlock expected, got {}".format(type(symbol))
        raise ValueError(msg)
    if isinstance(sym, list):
        sym = _sym.Group(sym)
    return sym, params
