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

import os
import ctypes
import mxnet as mx
from mxnet.base import SymbolHandle, check_call, _LIB, mx_uint, c_str_array, c_str
from mxnet.symbol import Symbol
import numpy as np
from mxnet.test_utils import assert_almost_equal


def _test_subgraph_exe(subgraph_backend):
    def _check_subgraph_exe1(sym, subgraph_backend, op_names):
        """Use the partitioned sym to simple_bind an executor and compare the outputs
        with those of the original executor"""
        out = SymbolHandle()
        check_call(_LIB.MXBuildSubgraphByOpNames(sym.handle, c_str(subgraph_backend), mx_uint(len(op_names)),
                                                  c_str_array(op_names), ctypes.byref(out)))

        partitioned_sym = Symbol(out)
        assert partitioned_sym.list_inputs() == sym.list_inputs()
        assert partitioned_sym.list_arguments() == sym.list_arguments()
        assert partitioned_sym.list_auxiliary_states() == sym.list_auxiliary_states()
        exe = sym.simple_bind(ctx=mx.current_context(), grad_req='null')
        partitioned_exe = partitioned_sym.simple_bind(ctx=mx.current_context(), grad_req='null')
        input_names = sym.list_inputs()
        for name in input_names:
            if name in exe.arg_dict:
                exe.arg_dict[name][:] = mx.nd.random.uniform(shape=exe.arg_dict[name].shape)
                partitioned_exe.arg_dict[name][:] = exe.arg_dict[name]
            else:
                assert name in exe.aux_dict
                exe.aux_dict[name][:] = mx.nd.random.uniform(shape=exe.aux_dict[name].shape)
                partitioned_exe.aux_dict[name][:] = exe.aux_dict[name]
        exe.forward()
        partitioned_exe.forward()
        assert len(exe.outputs) == len(partitioned_exe.outputs)
        for i in range(len(exe.outputs)):
            assert_almost_equal((exe.outputs[i] - partitioned_exe.outputs[i]).abs().sum().asnumpy(),
                                np.zeros(shape=(1,)))

    def _check_subgraph_exe2(sym, subgraph_backend, op_names):
        """Use env var MXNET_SUBGRAPH_BACKEND=default to trigger graph partitioning in simple_bind
        and compare results of the partitioned sym and the original sym."""
        def get_executor(sym, subgraph_backend=None, op_names=None, original_exec=None):
            if subgraph_backend is not None:
                os.environ['MXNET_SUBGRAPH_BACKEND'] = subgraph_backend
                check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)),
                                                             c_str_array(op_names)))
            exe = sym.simple_bind(ctx=mx.current_context(), grad_req='null')
            input_names = sym.list_inputs()
            for name in input_names:
                if name in exe.arg_dict:
                    exe.arg_dict[name][:] = mx.nd.random.uniform(shape=exe.arg_dict[name].shape)\
                        if original_exec is None else original_exec.arg_dict[name]
                else:
                    assert name in exe.aux_dict
                    exe.aux_dict[name][:] = mx.nd.random.uniform(shape=exe.aux_dict[name].shape)\
                        if original_exec is None else original_exec.aux_dict[name]
            exe.forward()
            if subgraph_backend is not None:
                check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))
                del os.environ['MXNET_SUBGRAPH_BACKEND']
            return exe

        original_exec = get_executor(sym)
        partitioned_exec = get_executor(sym, subgraph_backend, op_names, original_exec)
        outputs1 = original_exec.outputs
        outputs2 = partitioned_exec.outputs
        assert len(outputs1) == len(outputs2)
        for i in range(len(outputs1)):
            assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))

    def _check_subgraph_exe3(sym, subgraph_backend, op_names):
        """Use the partitioned sym to bind an executor and compare the outputs
        with those of the original executor"""
        out = SymbolHandle()
        check_call(_LIB.MXBuildSubgraphByOpNames(sym.handle, c_str(subgraph_backend), mx_uint(len(op_names)),
                                                  c_str_array(op_names), ctypes.byref(out)))

        partitioned_sym = Symbol(out)
        input_names = sym.list_inputs()
        arg_names = sym.list_arguments()
        aux_names = sym.list_auxiliary_states()
        assert partitioned_sym.list_inputs() == input_names
        assert partitioned_sym.list_arguments() == arg_names
        assert partitioned_sym.list_auxiliary_states() == aux_names
        arg_shapes, _, aux_shapes = sym.infer_shape()
        arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes]
        aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes]
        exe = sym.bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null')
        partitioned_exe = partitioned_sym.bind(ctx=mx.current_context(), args=arg_array,
                                               aux_states=aux_array, grad_req='null')
        exe.forward()
        partitioned_exe.forward()
        assert len(exe.outputs) == len(partitioned_exe.outputs)
        for i in range(len(exe.outputs)):
            assert_almost_equal((exe.outputs[i] - partitioned_exe.outputs[i]).abs().sum().asnumpy(),
                                np.zeros(shape=(1,)))

    def _check_subgraph_exe4(sym, subgraph_backend, op_names):
        """Use env var MXNET_SUBGRAPH_BACKEND=default to trigger graph partitioning in bind
        and compare results of the partitioned sym and the original sym."""
        def get_executor(sym, subgraph_backend=None, op_names=None, original_exec=None):
            if subgraph_backend is not None:
                os.environ['MXNET_SUBGRAPH_BACKEND'] = subgraph_backend
                check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)),
                                                             c_str_array(op_names)))
            arg_shapes, _, aux_shapes = sym.infer_shape()
            if subgraph_backend is None:
                arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes]
                aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes]
            else:
                arg_array = None
                aux_array = None
            exe = sym.bind(ctx=mx.current_context(),
                           args=arg_array if subgraph_backend is None else original_exec.arg_arrays,
                           aux_states=aux_array if subgraph_backend is None else original_exec.aux_arrays,
                           grad_req='null')
            exe.forward()
            if subgraph_backend is not None:
                check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))
                del os.environ['MXNET_SUBGRAPH_BACKEND']
            return exe

        original_exec = get_executor(sym)
        partitioned_exec = get_executor(sym, subgraph_backend, op_names, original_exec)
        outputs1 = original_exec.outputs
        outputs2 = partitioned_exec.outputs
        assert len(outputs1) == len(outputs2)
        for i in range(len(outputs1)):
            assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))

    def set_random_inputs(exe1, input_names):
        """Sets random values to exe1's args and auxs"""
        for name in input_names:
            if name in exe1.arg_dict:
                exe1.arg_dict[name][:] = mx.nd.random.uniform(shape=exe1.arg_dict[name].shape)
            else:
                assert name in exe1.aux_dict
                exe1.aux_dict[name][:] = mx.nd.random.uniform(shape=exe1.aux_dict[name].shape)

    def copy_inputs_between_executors(exe1, exe2, input_names):
        """Copies values of args and auxs from exe1 to exe2"""
        for name in input_names:
            if name in exe2.arg_dict:
                exe2.arg_dict[name][:] = exe1.arg_dict[name]
            else:
                assert name in exe2.aux_dict
                exe2.aux_dict[name][:] = exe1.aux_dict[name]

    def _check_subgraph_exe5(sym, subgraph_backend, op_names):
        """Call optimize_for to trigger graph partitioning without infer shapes/types before,
        then simple_bind and compare results of the partitioned sym and the original sym."""
        # simple_bind
        exe1 = sym.simple_bind(ctx=mx.current_context(), grad_req='null')
        input_names = sym.list_inputs()
        set_random_inputs(exe1, input_names)
        exe1.forward()

        # partition before simple_bind
        check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)),
                                                     c_str_array(op_names)))
        part_sym = sym.optimize_for(subgraph_backend)
        check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))

        exe2 = part_sym.simple_bind(ctx=mx.current_context(), grad_req='null')
        copy_inputs_between_executors(exe1, exe2, input_names)
        exe2.forward()

        # compare outputs
        outputs1 = exe1.outputs
        outputs2 = exe2.outputs
        assert len(outputs1) == len(outputs2)
        for i in range(len(outputs1)):
            assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))

    def _check_subgraph_exe6(sym, subgraph_backend, op_names):
        """Call optimize_for to trigger graph partitioning without infer shapes/types before,
        then simple_bind and compare results of the partitioned sym and the original sym."""
        # simple_bind
        exe1 = sym.simple_bind(ctx=mx.current_context(), grad_req='null')
        input_names = sym.list_inputs()
        set_random_inputs(exe1, input_names)
        exe1.forward()

        # infer shape/type before partition before simple_bind
        check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)),
                                                     c_str_array(op_names)))
        part_sym = sym.optimize_for(subgraph_backend, exe1.arg_dict)
        check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))

        exe2 = part_sym.simple_bind(ctx=mx.current_context(), grad_req='null')
        copy_inputs_between_executors(exe1, exe2, input_names)
        exe2.forward()

        # compare outputs
        outputs1 = exe1.outputs
        outputs2 = exe2.outputs
        assert len(outputs1) == len(outputs2)
        for i in range(len(outputs1)):
            assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))

    def _check_subgraph_exe7(sym, subgraph_backend, op_names):
        """Call optimize_for to trigger graph partitioning without infer shapes/types before,
        then bind and compare results of the partitioned sym and the original sym."""
        # bind
        arg_shapes, _, aux_shapes = sym.infer_shape()
        arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes]
        aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes]
        exe1 = sym.bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null')
        exe1.forward()

        # partition before bind
        check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)),
                                                     c_str_array(op_names)))
        part_sym = sym.optimize_for(subgraph_backend)
        check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))

        exe2 = part_sym.bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null')
        exe2.forward()

        # compare outputs
        outputs1 = exe1.outputs
        outputs2 = exe2.outputs
        assert len(outputs1) == len(outputs2)
        for i in range(len(outputs1)):
            assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))

    def _check_subgraph_exe8(sym, subgraph_backend, op_names):
        """Call optimize_for to infer shapes, types and dtypes followed by graph partitioning,
        then bind and compare results of the partitioned sym and the original sym."""
        # bind
        arg_shapes, _, aux_shapes = sym.infer_shape()
        arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes]
        aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes]
        exe1 = sym.bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null')
        exe1.forward()

        # infer shape/type before partition before bind
        check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)),
                                                     c_str_array(op_names)))
        part_sym = sym.optimize_for(subgraph_backend, arg_array)
        check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))

        exe2 = part_sym.bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null')
        exe2.forward()

        # compare outputs
        outputs1 = exe1.outputs
        outputs2 = exe2.outputs
        assert len(outputs1) == len(outputs2)
        for i in range(len(outputs1)):
            assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), np.zeros(shape=(1,)))

    def check_subgraph_exe(sym, subgraph_backend, op_names):
        _check_subgraph_exe1(sym, subgraph_backend, op_names)
        _check_subgraph_exe2(sym, subgraph_backend, op_names)
        _check_subgraph_exe3(sym, subgraph_backend, op_names)
        _check_subgraph_exe4(sym, subgraph_backend, op_names)
        _check_subgraph_exe5(sym, subgraph_backend, op_names)
        _check_subgraph_exe6(sym, subgraph_backend, op_names)
        _check_subgraph_exe7(sym, subgraph_backend, op_names)
        _check_subgraph_exe8(sym, subgraph_backend, op_names)

    def test_network_structure_1(subgraph_backend):
        data1 = mx.sym.var('data1', shape=(2, 3, 10, 10))
        data2 = mx.sym.var('data2')
        conv1 = mx.sym.Convolution(data=data1, weight=data2, no_bias=True, kernel=(2, 2), num_filter=1)
        conv2 = mx.sym.Convolution(data=data2, no_bias=True, kernel=(1, 1), num_filter=1)
        out = mx.sym.Group([conv1, conv2])
        check_subgraph_exe(out, subgraph_backend, ['Convolution'])

    def test_network_structure_2(subgraph_backend):
        # this tests whether the partitioning algorithm can deal with cycles
        data = mx.sym.var('data', shape=(2, 3, 10, 10))
        ret = mx.sym.exp(data)
        ret1 = mx.sym.cos(ret)
        ret2 = mx.sym.sin(ret)
        ret = ret1 + ret2
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'sin', '_Plus', 'elemwise_add', '_plus'])
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'cos', '_Plus', 'elemwise_add', '_plus'])

    def test_network_structure_3(subgraph_backend):
        # this tests whether the partitioned sym can distinguish in_args and aux_states
        data = mx.sym.var('data', shape=(2, 3, 10, 10))
        ret = mx.sym.exp(data)
        ret1 = mx.sym.cos(ret)
        ret2 = mx.sym.sin(ret)
        ret = ret1 + ret2
        ret = mx.sym.BatchNorm(ret)
        ret = mx.sym.BatchNorm(ret)
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'sin', '_Plus', 'elemwise_add', '_plus'])
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'cos', '_Plus', 'elemwise_add', '_plus'])
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'sin', '_Plus', 'elemwise_add', '_plus', 'BatchNorm'])
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'cos', '_Plus', 'elemwise_add', '_plus', 'BatchNorm'])
        check_subgraph_exe(ret, subgraph_backend, ['exp', 'BatchNorm'])
        check_subgraph_exe(ret, subgraph_backend, ['BatchNorm'])

    def test_network_structure_4(subgraph_backend):
        # the last op has multiple duplicate outputs
        data = mx.sym.var('data', shape=(2, 3, 10, 10))
        ret = mx.sym.exp(data)
        ret = mx.sym.Group([ret, ret, ret])
        check_subgraph_exe(ret, subgraph_backend, ['exp'])

    def test_network_structure_5(subgraph_backend):
        # the subgraph has two duplicate input entries
        data = mx.sym.var('data', shape=(2, 3, 10, 10))
        ret = data + data
        check_subgraph_exe(ret, subgraph_backend, ['_plus', '_Plus', 'elemwise_add'])

    def test_network_structure_6(subgraph_backend):
        def get_graph():
            data1 = mx.sym.Variable('data1', shape=(3, 3, 10, 10), dtype=np.float32)
            data2 = mx.sym.Variable('data2', shape=(1, 0, 2, 2))
            data3 = mx.sym.sin(data2)
            conv = mx.sym.Convolution(data=data1, weight=data3, kernel=(2, 2), num_filter=1)
            rets = [(conv, []),
                    (conv, [mx.sym.sin.__name__]),
                    (conv, [mx.sym.Convolution.__name__]),
                    (conv, [mx.sym.sin.__name__, mx.sym.Convolution.__name__])]
            return rets

        for sym, op_names in get_graph():
            check_subgraph_exe(sym, subgraph_backend, op_names)

    def test_network_structure_7(subgraph_backend):
        # in this graph, the subgraph node and the other two external nodes form a cycle
        data = mx.sym.Variable('data', shape=(1,))
        ret1 = mx.sym.sin(data)
        ret2 = mx.sym.cos(ret1)
        for _ in range(5):
            ret2 = mx.sym.cos(ret2)
        ret = ret1 + ret2
        check_subgraph_exe(ret, subgraph_backend, ['sin', 'elemwise_add', '_plus', '_Plus'])

    test_network_structure_1(subgraph_backend)
    test_network_structure_2(subgraph_backend)
    test_network_structure_3(subgraph_backend)
    test_network_structure_4(subgraph_backend)
    test_network_structure_5(subgraph_backend)
    test_network_structure_6(subgraph_backend)
    test_network_structure_7(subgraph_backend)

def test_subgraph_exe():
    _test_subgraph_exe('default')

def test_subgraph_v2_exe():
    _test_subgraph_exe('default_v2')

if __name__ == '__main__':
    import nose
    nose.runmodule()