# 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.

"""
Tests for individual operators
This module contains operator tests which currently do not exist on
ONNX backend test framework. Once we have PRs on the ONNX repo and get
those PRs merged, this file will get EOL'ed.
"""
# pylint: disable=too-many-locals,wrong-import-position,import-error
from __future__ import absolute_import
import sys
import os
import unittest
import logging
import tarfile
from collections import namedtuple
import numpy as np
import numpy.testing as npt
from onnx import checker, numpy_helper, helper, load_model
from onnx import TensorProto
from mxnet.test_utils import download
from mxnet.contrib import onnx as onnx_mxnet
import mxnet as mx
import backend

CURR_PATH = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
sys.path.insert(0, os.path.join(CURR_PATH, '../../python/unittest'))

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)


def get_rnd(shape, low=-1.0, high=1.0, dtype=np.float32):
    if dtype == np.float32:
        return (np.random.uniform(low, high,
                                  np.prod(shape)).reshape(shape).astype(np.float32))
    elif dtype == np.int32:
        return (np.random.randint(low, high,
                                  np.prod(shape)).reshape(shape).astype(np.float32))
    elif dtype == np.bool_:
        return np.random.choice(a=[False, True], size=shape).astype(np.float32)


def _fix_attributes(attrs, attribute_mapping):
    new_attrs = attrs
    attr_modify = attribute_mapping.get('modify', {})
    for k, v in attr_modify.items():
        new_attrs[v] = new_attrs.pop(k, None)

    attr_add = attribute_mapping.get('add', {})
    for k, v in attr_add.items():
        new_attrs[k] = v

    attr_remove = attribute_mapping.get('remove', [])
    for k in attr_remove:
        if k in new_attrs:
            del new_attrs[k]

    return new_attrs


def forward_pass(sym, arg, aux, data_names, input_data):
    """ Perform forward pass on given data
    :param sym: Symbol
    :param arg: Arg params
    :param aux: Aux params
    :param data_names: Input names (list)
    :param input_data: Input data (list). If there is only one input,
                        pass it as a list. For example, if input is [1, 2],
                        pass input_data=[[1, 2]]
    :return: result of forward pass
    """
    data_shapes = []
    data_forward = []
    for idx in range(len(data_names)):
        val = input_data[idx]
        data_shapes.append((data_names[idx], np.shape(val)))
        data_forward.append(mx.nd.array(val))
    # create module
    mod = mx.mod.Module(symbol=sym, data_names=data_names, context=mx.cpu(), label_names=None)
    mod.bind(for_training=False, data_shapes=data_shapes, label_shapes=None)
    if not arg and not aux:
        mod.init_params()
    else:
        mod.set_params(arg_params=arg, aux_params=aux,
                       allow_missing=True, allow_extra=True)
    # run inference
    batch = namedtuple('Batch', ['data'])
    mod.forward(batch(data_forward), is_train=False)

    return mod.get_outputs()[0].asnumpy()


def get_input_tensors(input_data):
    input_tensor = []
    input_names = []
    input_sym = []
    for idx, ip in enumerate(input_data):
        name = "input" + str(idx + 1)
        input_sym.append(mx.sym.Variable(name))
        input_names.append(name)
        input_tensor.append(helper.make_tensor_value_info(name,
                                                          TensorProto.FLOAT, shape=np.shape(ip)))
    return input_names, input_tensor, input_sym


def get_onnx_graph(testname, input_names, inputs, output_name, output_shape, attr):
    outputs = [helper.make_tensor_value_info("output", TensorProto.FLOAT, shape=output_shape)]

    nodes = [helper.make_node(output_name, input_names, ["output"], **attr)]

    graph = helper.make_graph(nodes, testname, inputs, outputs)

    model = helper.make_model(graph)
    return model

class TestNode(unittest.TestCase):
    """ Tests for models.
    Tests are dynamically added.
    Therefore edit test_models to add more tests.
    """
    def test_import_export(self):
        for test in test_cases:
            test_name, mxnet_op, onnx_name, inputs, attrs, mxnet_specific, fix_attrs, check_value, check_shape = test
            with self.subTest(test_name):
                names, input_tensors, inputsym = get_input_tensors(inputs)
                if inputs:
                    test_op = mxnet_op(*inputsym, **attrs)
                    mxnet_output = forward_pass(test_op, None, None, names, inputs)
                    outputshape = np.shape(mxnet_output)
                else:
                    test_op = mxnet_op(**attrs)
                    shape = attrs.get('shape', (1,))
                    x = mx.nd.zeros(shape, dtype='float32')
                    xgrad = mx.nd.zeros(shape, dtype='float32')
                    exe = test_op.bind(ctx=mx.cpu(), args={'x': x}, args_grad={'x': xgrad})
                    mxnet_output = exe.forward(is_train=False)[0].asnumpy()
                    outputshape = np.shape(mxnet_output)

                if mxnet_specific:
                    onnxmodelfile = onnx_mxnet.export_model(test_op, {}, [np.shape(ip) for ip in inputs],
                                                            [ip.dtype for ip in inputs],
                                                            onnx_name + ".onnx")
                    onnxmodel = load_model(onnxmodelfile)
                else:
                    onnx_attrs = _fix_attributes(attrs, fix_attrs)
                    onnxmodel = get_onnx_graph(test_name, names, input_tensors, onnx_name, outputshape, onnx_attrs)

                bkd_rep = backend.prepare(onnxmodel, operation='export')
                output = bkd_rep.run(inputs)

                if check_value:
                    npt.assert_almost_equal(output[0], mxnet_output)

                if check_shape:
                    npt.assert_equal(output[0].shape, outputshape)

        input1 = get_rnd((1, 10, 2, 3))
        ipsym = mx.sym.Variable("input1")
        for test in test_scalar_ops:
            if test == 'Add':
                outsym = 2 + ipsym
            if test == "Sub":
                outsym = ipsym - 2
            if test == "rSub":
                outsym = ipsym.__rsub__(2)
            if test == "Mul":
                outsym = 2 * ipsym
            if test == "Div":
                outsym = ipsym / 2
            if test == "Pow":
                outsym = ipsym ** 2
            forward_op = forward_pass(outsym, None, None, ['input1'], input1)
            converted_model = onnx_mxnet.export_model(outsym, {}, [np.shape(input1)], np.float32,
                                                      onnx_file_path=outsym.name + ".onnx")

            sym, arg_params, aux_params = onnx_mxnet.import_model(converted_model)
            result = forward_pass(sym, arg_params, aux_params, ['input1'], input1)

            npt.assert_almost_equal(result, forward_op)

    def test_imports(self):
        for test in import_test_cases:
            test_name, onnx_name, inputs, np_op, attrs = test
            with self.subTest(test_name):
                names, input_tensors, inputsym = get_input_tensors(inputs)
                np_out = [np_op(*inputs, **attrs)]
                output_shape = np.shape(np_out)
                onnx_model = get_onnx_graph(test_name, names, input_tensors, onnx_name, output_shape, attrs)
                bkd_rep = backend.prepare(onnx_model, operation='import')
                mxnet_out = bkd_rep.run(inputs)
                npt.assert_almost_equal(np_out, mxnet_out, decimal=4)

    def test_exports(self):
        input_shape = (2,1,3,1)
        for test in export_test_cases:
            test_name, onnx_name, mx_op, attrs = test
            input_sym = mx.sym.var('data')
            outsym = mx_op(input_sym, **attrs)
            converted_model = onnx_mxnet.export_model(outsym, {}, [input_shape], np.float32,
                                                      onnx_file_path=outsym.name + ".onnx")
            model = load_model(converted_model)
            checker.check_model(model)


# test_case = ("test_case_name", mxnet op, "ONNX_op_name", [input_list], attribute map, MXNet_specific=True/False,
# fix_attributes = {'modify': {mxnet_attr_name: onnx_attr_name},
#                   'remove': [attr_name],
#                   'add': {attr_name: value},
# check_value=True/False, check_shape=True/False)
test_cases = [
    ("test_equal", mx.sym.broadcast_equal, "Equal", [get_rnd((1, 3, 4, 5)), get_rnd((1, 5))], {}, False, {}, True,
     False),
    ("test_greater", mx.sym.broadcast_greater, "Greater", [get_rnd((1, 3, 4, 5)), get_rnd((1, 5))], {}, False, {}, True,
     False),
    ("test_less", mx.sym.broadcast_lesser, "Less", [get_rnd((1, 3, 4, 5)), get_rnd((1, 5))], {}, False, {}, True,
     False),
    ("test_and", mx.sym.broadcast_logical_and, "And",
     [get_rnd((3, 4, 5), dtype=np.bool_), get_rnd((3, 4, 5), dtype=np.bool_)], {}, False, {}, True, False),
    ("test_xor", mx.sym.broadcast_logical_xor, "Xor",
     [get_rnd((3, 4, 5), dtype=np.bool_), get_rnd((3, 4, 5), dtype=np.bool_)], {}, False, {}, True, False),
    ("test_or", mx.sym.broadcast_logical_or, "Or",
     [get_rnd((3, 4, 5), dtype=np.bool_), get_rnd((3, 4, 5), dtype=np.bool_)], {}, False, {}, True, False),
    ("test_not", mx.sym.logical_not, "Not", [get_rnd((3, 4, 5), dtype=np.bool_)], {}, False, {}, True, False),
    ("test_square", mx.sym.square, "Pow", [get_rnd((2, 3), dtype=np.int32)], {}, True, {}, True, False),
    ("test_spacetodepth", mx.sym.space_to_depth, "SpaceToDepth", [get_rnd((1, 1, 4, 6))],
     {'block_size': 2}, False, {}, True, False),
    ("test_softmax", mx.sym.SoftmaxOutput, "Softmax", [get_rnd((1000, 1000)), get_rnd(1000)],
     {'ignore_label': 0, 'use_ignore': False}, True, {}, True, False),
    ("test_logistic_regression", mx.sym.LogisticRegressionOutput, "Sigmoid",
     [get_rnd((1000, 1000)), get_rnd((1000, 1000))], {}, True, {}, True, False),
    # TODO: After rewrite, FC would fail this testcase. Commenting this out for now
    # ("test_fullyconnected", mx.sym.FullyConnected, "Gemm", [get_rnd((4, 3)), get_rnd((4, 3)), get_rnd(4)],
    #  {'num_hidden': 4, 'name': 'FC'}, True, {}, True, False),
    ("test_lppool1", mx.sym.Pooling, "LpPool", [get_rnd((2, 3, 20, 20))],
     {'kernel': (4, 5), 'pad': (0, 0), 'stride': (1, 1), 'p_value': 1, 'pool_type': 'lp'}, False,
     {'modify': {'kernel': 'kernel_shape', 'pad': 'pads', 'stride': 'strides', 'p_value': 'p'},
      'remove': ['pool_type']}, True, False),
    ("test_lppool2", mx.sym.Pooling, "LpPool", [get_rnd((2, 3, 20, 20))],
     {'kernel': (4, 5), 'pad': (0, 0), 'stride': (1, 1), 'p_value': 2, 'pool_type': 'lp'}, False,
     {'modify': {'kernel': 'kernel_shape', 'pad': 'pads', 'stride': 'strides', 'p_value': 'p'},
      'remove': ['pool_type']}, True, False),
    ("test_globallppool1", mx.sym.Pooling, "GlobalLpPool", [get_rnd((2, 3, 20, 20))],
     {'kernel': (4, 5), 'pad': (0, 0), 'stride': (1, 1), 'p_value': 1, 'pool_type': 'lp', 'global_pool': True}, False,
     {'modify': {'p_value': 'p'},
      'remove': ['pool_type', 'kernel', 'pad', 'stride', 'global_pool']}, True, False),
    ("test_globallppool2", mx.sym.Pooling, "GlobalLpPool", [get_rnd((2, 3, 20, 20))],
     {'kernel': (4, 5), 'pad': (0, 0), 'stride': (1, 1), 'p_value': 2, 'pool_type': 'lp', 'global_pool': True}, False,
     {'modify': {'p_value': 'p'},
      'remove': ['pool_type', 'kernel', 'pad', 'stride', 'global_pool']}, True, False),
    ("test_roipool", mx.sym.ROIPooling, "MaxRoiPool",
     [[[get_rnd(shape=(8, 6), low=1, high=100, dtype=np.int32)]], [[0, 0, 0, 4, 4]]],
     {'pooled_size': (2, 2), 'spatial_scale': 0.7}, False,
     {'modify': {'pooled_size': 'pooled_shape'}}, True, False),

    # since results would be random, checking for shape alone
    ("test_multinomial", mx.sym.sample_multinomial, "Multinomial",
     [np.array([0, 0.1, 0.2, 0.3, 0.4]).astype("float32")],
     {'shape': (10,)}, False, {'modify': {'shape': 'sample_size'}}, False, True),
    ("test_random_normal", mx.sym.random_normal, "RandomNormal", [],
     {'shape': (2, 2), 'loc': 0, 'scale': 1}, False, {'modify': {'loc': 'mean'}}, False, True),
    ("test_random_uniform", mx.sym.random_uniform, "RandomUniform", [],
     {'shape': (2, 2), 'low': 0.5, 'high': 1.0}, False, {}, False, True)
]

test_scalar_ops = ['Add', 'Sub', 'rSub' 'Mul', 'Div', 'Pow']

# test_case = ("test_case_name", "ONNX_op_name", [input_list], np_op, attribute map)
import_test_cases = [
    ("test_lpnormalization_default", "LpNormalization", [get_rnd([5, 3, 3, 2])], np.linalg.norm, {'ord':2, 'axis':-1}),
    ("test_lpnormalization_ord1", "LpNormalization", [get_rnd([5, 3, 3, 2])], np.linalg.norm, {'ord':1, 'axis':-1}),
    ("test_lpnormalization_ord2", "LpNormalization", [get_rnd([5, 3, 3, 2])], np.linalg.norm, {'ord':2, 'axis':1})
]

# test_case = ("test_case_name", "ONNX_op_name", mxnet_op, attribute map)
export_test_cases = [
    ("test_expand", "Expand", mx.sym.broadcast_to, {'shape': (2,1,3,1)}),
    ("test_tile", "Tile", mx.sym.tile, {'reps': (2,3)})
]

if __name__ == '__main__':
    unittest.main()