# 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 mxnet as mx
import mxnet.ndarray as nd
from mxnet.test_utils import *
import numpy as np
from functools import reduce
from mxnet.module.executor_group import DataParallelExecutorGroup
from common import setup_module, with_seed, assertRaises, teardown
from collections import namedtuple


@with_seed()
def test_module_dtype():
    dtype = np.float16
    dshape = (3, 8, 7)

    sym = mx.sym.Variable('data')
    sym = mx.sym.Activation(data=sym, act_type='relu', __layout__='TNC')

    mod = mx.mod.Module(sym, ('data',), None, context=[mx.cpu(0), mx.cpu(1)])
    mod.bind(data_shapes=[mx.io.DataDesc('data', dshape, dtype, layout='TNC')])
    mod.init_params()
    mod.forward(mx.io.DataBatch(data=[mx.nd.ones(dshape, dtype=dtype)],
                              label=None))
    mod.backward([mx.nd.ones(dshape, dtype=dtype)])

    for x in mod.get_outputs():
      assert x.dtype == dtype


def test_module_bind():
    sym = mx.sym.Variable('data')
    sym = mx.sym.Activation(data=sym, act_type='relu', __layout__='TNC')

    mod = mx.mod.Module(sym, ('data',), None, context=[mx.cpu(0), mx.cpu(1)])
    assertRaises(TypeError, mod.bind, data_shapes=[('data', mx.nd.array([10,10]))])
    assert mod.binded == False

    mod.bind(data_shapes=[('data', (10,10))])
    assert mod.binded == True


@with_seed()
def test_module_input_grads():
    a = mx.sym.Variable('a', __layout__='NC')
    b = mx.sym.Variable('b', __layout__='NC')
    c = mx.sym.Variable('c', __layout__='NC')

    c = a + 2 * b + 3 * c
    net = mx.mod.Module(c, data_names=['b', 'c', 'a'], label_names=None,
                        context=[mx.cpu(0), mx.cpu(1)])
    net.bind(data_shapes=[['b', (5, 5)], ['c', (5, 5)], ['a', (5, 5)]],
             label_shapes=None, inputs_need_grad=True)
    net.init_params()

    net.forward(data_batch=mx.io.DataBatch(data=[nd.ones((5, 5)),
                                                 nd.ones((5, 5)),
                                                 nd.ones((5, 5))]))
    net.backward(out_grads=[nd.ones((5, 5))])
    input_grads = net.get_input_grads()
    b_grad = input_grads[0].asnumpy()
    c_grad = input_grads[1].asnumpy()
    a_grad = input_grads[2].asnumpy()
    assert np.all(a_grad == 1), a_grad
    assert np.all(b_grad == 2), b_grad
    assert np.all(c_grad == 3), c_grad


@with_seed()
def test_module_ctx_group():
    def check_module_ctx_group(ctxs, group2ctxs, grad_ctxs=None):
        with mx.AttrScope(ctx_group='dev1'):
            a = mx.symbol.Variable('a')
            a = a * 2
        with mx.AttrScope(ctx_group='dev2'):
            b = mx.symbol.Variable('b')
            c = a + b
        shape = (2, 5)
        mod1 = mx.mod.Module(c, context=ctxs, data_names=['a', 'b'], label_names=None,
                             group2ctxs=group2ctxs)
        mod1.bind(data_shapes=[['a', shape], ['b', shape]], inputs_need_grad=True)
        mod1.init_params()
        mod1.forward(data_batch=mx.io.DataBatch(data=[mx.nd.ones(shape), mx.nd.ones(shape)]), is_train=True)
        mod1.backward([mx.nd.ones(shape)])
        mod1_input_grads = mod1.get_input_grads()

        mod2 = mx.mod.Module(c, context=ctxs, data_names=['a', 'b'], label_names=None)
        mod2.bind(data_shapes=[['a', shape], ['b', shape]], inputs_need_grad=True)
        mod2.init_params()
        mod2.forward(data_batch=mx.io.DataBatch(data=[mx.nd.ones(shape), mx.nd.ones(shape)]), is_train=True)
        mod2.backward([mx.nd.ones(shape)])
        mod2_input_grads = mod2.get_input_grads()

        if grad_ctxs is not None:
            assert(mod1_input_grads[0].context == grad_ctxs[0])
            assert(mod1_input_grads[1].context == grad_ctxs[1])
        assert(np.all(mod1_input_grads[0].asnumpy() == mod2_input_grads[0].asnumpy()))
        assert(np.all(mod1_input_grads[1].asnumpy() == mod2_input_grads[1].asnumpy()))

    check_module_ctx_group([mx.cpu(0)], {'dev1': mx.cpu(1), 'dev2': mx.cpu(2)}, grad_ctxs=[mx.cpu(1), mx.cpu(2)])
    check_module_ctx_group([mx.cpu(0), mx.cpu(1)],
        [{'dev1': mx.cpu(2), 'dev2': mx.cpu(3)}, {'dev1': mx.cpu(4), 'dev2': mx.cpu(5)}])
    check_module_ctx_group([mx.cpu(0), mx.cpu(1)], {'dev1': mx.cpu(2), 'dev2': mx.cpu(3)})
    check_module_ctx_group([mx.cpu(0), mx.cpu(1)], {'dev1': mx.cpu(2), 'dev2': [mx.cpu(3)]})
    check_module_ctx_group([mx.cpu(0), mx.cpu(1)], {'dev1':mx.cpu(2), 'dev2':[mx.cpu(3), mx.cpu(3)]})
    check_module_ctx_group([mx.cpu(0), mx.cpu(1)],
        {'dev1':[mx.cpu(2), mx.cpu(2)], 'dev2':[mx.cpu(3), mx.cpu(3)]})

@with_seed()
def test_bucket_module_ctx_group():
    num_hidden = 10
    batch_size = 5
    def sym_gen(seq_len):
        with mx.AttrScope(ctx_group='dev1'):
            data = mx.symbol.Variable('data')
            weight = mx.symbol.Variable('dev1_weight')
            bias = mx.symbol.Variable('dev1_bias')
            fc = data
            for i in range(seq_len):
                fc  = mx.symbol.FullyConnected(data=fc, weight=weight, bias=bias,
                                               name='dev1_fc_%d' % i, num_hidden=num_hidden)
        with mx.AttrScope(ctx_group='dev2'):
            label = mx.symbol.Variable('label')
            weight = mx.symbol.Variable('dev2_weight')
            bias = mx.symbol.Variable('dev2_bias')
            for i in range(seq_len):
                fc  = mx.symbol.FullyConnected(data=fc, weight=weight, bias=bias,
                                               name='dev2_fc_%d' % i, num_hidden=num_hidden)
            sym = mx.symbol.SoftmaxOutput(fc, label, name='softmax')

        return sym, ('data',), ('label',)

    mod = mx.mod.BucketingModule(sym_gen=sym_gen, default_bucket_key=10, context=[mx.cpu(0)],
                                 group2ctxs=[{'dev1': mx.cpu(1), 'dev2': mx.cpu(2)}])
    mod.bind(data_shapes=[['data', (batch_size, num_hidden)]],
             label_shapes=[['label', (batch_size,)]],
             for_training=True, inputs_need_grad=True)
    assert(mod.binded)

@with_seed()
def test_module_layout():
    sym = mx.sym.Variable('data')
    sym = mx.sym.Activation(data=sym, act_type='relu', __layout__='TNC')

    dshape = (3, 8, 7)
    mod = mx.mod.Module(sym, ('data',), None, context=[mx.cpu(0), mx.cpu(1)])
    mod.bind(data_shapes=[mx.io.DataDesc('data', dshape, layout='TNC')])
    mod.init_params()
    mod.forward(mx.io.DataBatch(data=[mx.nd.ones(dshape)],
                                label=None))
    mod.backward([mx.nd.ones(dshape)])
    assert mod.get_outputs()[0].shape == dshape

    hdshape = (3, 4, 7)
    for x in mod.get_outputs(merge_multi_context=False)[0]:
        assert x.shape == hdshape


@with_seed()
def test_save_load():
    previous_update_on_kvstore = os.getenv('MXNET_UPDATE_ON_KVSTORE', "1")
    os.putenv('MXNET_UPDATE_ON_KVSTORE', '1')
    def dict_equ(a, b):
        assert set(a) == set(b)
        for k in a:
            assert (a[k].asnumpy() == b[k].asnumpy()).all()

    sym = mx.sym.Variable('data')
    sym = mx.sym.FullyConnected(sym, num_hidden=100)

    # single device
    mod = mx.mod.Module(sym, ('data',))
    mod.bind(data_shapes=[('data', (10, 10))])
    mod.init_params()
    mod.init_optimizer(optimizer_params={'learning_rate':0.1, 'momentum':0.9})
    mod.update()
    mod.save_checkpoint('test', 0, save_optimizer_states=True)

    mod2 = mx.mod.Module.load('test', 0, load_optimizer_states=True, data_names=('data',))
    mod2.bind(data_shapes=[('data', (10, 10))])
    mod2.init_optimizer(optimizer_params={'learning_rate':0.1, 'momentum':0.9})
    assert mod._symbol.tojson() == mod2._symbol.tojson()
    dict_equ(mod.get_params()[0], mod2.get_params()[0])
    dict_equ(mod._updater.states, mod2._updater.states)

    # multi device
    mod = mx.mod.Module(sym, ('data',), context=[mx.cpu(0), mx.cpu(1)])
    mod.bind(data_shapes=[('data', (10, 10))])
    mod.init_params()
    mod.init_optimizer(optimizer_params={'learning_rate':0.1, 'momentum':0.9})
    mod.update()
    mod.save_checkpoint('test', 0, save_optimizer_states=True)

    mod2 = mx.mod.Module.load('test', 0, load_optimizer_states=True, data_names=('data',))
    mod2.bind(data_shapes=[('data', (10, 10))])
    mod2.init_optimizer(optimizer_params={'learning_rate':0.1, 'momentum':0.9})
    assert mod._symbol.tojson() == mod2._symbol.tojson()
    dict_equ(mod.get_params()[0], mod2.get_params()[0])
    dict_equ(mod._kvstore._updater.states, mod2._updater.states)
    os.putenv('MXNET_UPDATE_ON_KVSTORE', previous_update_on_kvstore)


@with_seed()
def test_module_reshape():
    data = mx.sym.Variable('data')
    sym = mx.sym.FullyConnected(data, num_hidden=20, name='fc')

    dshape = (7, 20)
    mod = mx.mod.Module(sym, ('data',), None, context=[mx.cpu(0), mx.cpu(1)])
    mod.bind(data_shapes=[('data', dshape)])
    mod.init_params()
    mod.init_optimizer(optimizer_params={'learning_rate': 1})

    mod.forward(mx.io.DataBatch(data=[mx.nd.ones(dshape)],
                                label=None))
    mod.backward([mx.nd.ones(dshape)])
    mod.update()
    assert mod.get_outputs()[0].shape == dshape
    assert (mod.get_params()[0]['fc_bias'].asnumpy() == -1).all()

    dshape = (14, 20)
    mod.reshape(data_shapes=[('data', dshape)])
    mod.forward(mx.io.DataBatch(data=[mx.nd.ones(dshape)],
                                label=None))
    mod.backward([mx.nd.ones(dshape)])
    mod.update()
    assert mod.get_outputs()[0].shape == dshape
    assert (mod.get_params()[0]['fc_bias'].asnumpy() == -3).all()


@with_seed()
def test_module_states():
    stack = mx.rnn.SequentialRNNCell()
    for i in range(2):
        stack.add(mx.rnn.LSTMCell(num_hidden=20, prefix='lstm_l%d_'%i))
    begin_state = stack.begin_state(func=mx.sym.Variable)
    _, states = stack.unroll(10, begin_state=begin_state, inputs=mx.sym.Variable('data'))

    state_names = [i.name for i in begin_state]
    mod = mx.mod.Module(mx.sym.Group(states), context=[mx.cpu(0), mx.cpu(1)],
                        label_names=None, state_names=state_names)
    mod.bind(data_shapes=[('data', (5, 10))], label_shapes=None, for_training=False)
    mod.init_params()
    batch = mx.io.DataBatch(data=[mx.nd.zeros((5, 10))], label=[])

    mod.set_states(value=1)
    mod.forward(batch)
    out = mod.get_outputs(merge_multi_context=False)
    out1 = mod.get_outputs(merge_multi_context=True)

    mod.set_states(states=out)
    mod.forward(batch)
    out2 = mod.get_outputs(merge_multi_context=True)

    for x1, x2 in zip(out1, out2):
        assert not mx.test_utils.almost_equal(x1.asnumpy(), x2.asnumpy(), rtol=1e-3)


@with_seed()
def test_module_switch_bucket():
    vocab_dim = 5000
    num_hidden = 100
    num_embedding = 100
    num_layer = 2
    default_key = 10
    test_key = 5
    batch_size = 32
    contexts = [mx.cpu(0)]
    initializer = mx.init.Xavier(factor_type="in", magnitude=2.34)

    #generate symbols for an LSTM network
    def sym_gen(seq_len):
        data = mx.sym.Variable('data')
        label = mx.sym.Variable('softmax_label')
        embed = mx.sym.Embedding(data=data, input_dim=vocab_dim,
                                 output_dim=num_embedding)
        stack = mx.rnn.SequentialRNNCell()
        for i in range(num_layer):
            stack.add(mx.rnn.LSTMCell(num_hidden=num_hidden, prefix='lstm_l%d_'%i))
        outputs, states = stack.unroll(seq_len, inputs=embed, merge_outputs=True)

        pred = mx.sym.Reshape(outputs, shape=(-1, num_hidden))
        pred = mx.sym.FullyConnected(data=pred, num_hidden=vocab_dim, name='pred')

        label = mx.sym.Reshape(label, shape=(-1,))
        pred = mx.sym.SoftmaxOutput(data=pred, label=label, name='softmax')

        return pred, ('data',), ('softmax_label',)

    def create_bucketing_module(key):
        model = mx.mod.BucketingModule(
            sym_gen             = sym_gen,
            default_bucket_key  = key,
            context             = contexts)
        model.bind([('data', (batch_size, key))],
                    [('softmax_label', (batch_size, key))], True, False)
        model.init_params(initializer=initializer)
        return model
    #initialize the bucketing module with the default bucket key
    bucketing_model = create_bucketing_module(default_key)
    #check name
    assert bucketing_model.symbol.list_arguments()[1] == "embedding0_weight",\
        "Error in assigning names for args in BucketingModule"

    #switch to test_key
    bucketing_model.switch_bucket(test_key, [('data', (batch_size, test_key))],
                                  [('softmax_label', (batch_size, test_key))])
    total_bytes_before = bucketing_model._buckets[default_key]._total_exec_bytes

    #remove test_key and switch again
    del bucketing_model._buckets[test_key]
    bucketing_model.switch_bucket(test_key, [('data', (batch_size, test_key))],
                                  [('softmax_label', (batch_size, test_key))])
    total_bytes_after = bucketing_model._buckets[default_key]._total_exec_bytes
    #the default bucket is expected to reuse the bytes allocated
    assert total_bytes_after == total_bytes_before


# roywei: Getting rid of fixed seed as flakiness could not be reproduced,
# tracked at: https://github.com/apache/incubator-mxnet/issues/11705
@with_seed()
def test_module_set_params():
    # data iter
    data = mx.nd.array([[0.05, .10]]);
    label = mx.nd.array([[.01, 0.99]]);
    train_data = mx.io.NDArrayIter(data, label, batch_size=1)

    # symbols
    x = mx.symbol.Variable('data')
    x = mx.symbol.FullyConnected(name='fc_0', data=x, num_hidden=2)
    x = mx.symbol.Activation(name="act_0", data=x, act_type='sigmoid')
    x = mx.symbol.FullyConnected(name='fc_1', data=x, num_hidden=2)
    x = mx.symbol.Activation(name="act_1", data=x, act_type='sigmoid')
    x = mx.symbol.LinearRegressionOutput(data=x, name='softmax', grad_scale=2)

    # create module
    mod = mx.mod.Module(x, context=[mx.cpu()]);
    mod.bind(train_data.provide_data, label_shapes=train_data.provide_label,
             for_training=True)

    arg_params_correct = {'fc_0_weight': mx.nd.array([[.15, .20], [.25, .30]]),
                  'fc_0_bias'  : mx.nd.array([.35, .35]),
                  'fc_1_weight': mx.nd.array([[.40, .45], [.50, .55]]),
                  'fc_1_bias'  : mx.nd.array([.60, .60])}

    arg_params_missing = {'fc_0_weight': mx.nd.array([[.15, .20], [.25, .30]]),
                  'fc_0_bias'  : mx.nd.array([.35, .35]),
                  'fc_1_weight': mx.nd.array([[.40, .45], [.50, .55]])}

    arg_params_extra = {'fc_0_weight': mx.nd.array([[.15, .20], [.25, .30]]),
                  'fc_0_bias'  : mx.nd.array([.35, .35]),
                  'fc_1_weight': mx.nd.array([[.40, .45], [.50, .55]]),
                  'fc_1_bias'  : mx.nd.array([.60, .60]),
                  'fc_2_weight': mx.nd.array([.60, .60])}

    arg_params_missing_extra = {'fc_2_weight': mx.nd.array([.60, .60])}

    # test regular set_params
    mod.set_params(force_init=True, arg_params=arg_params_correct, aux_params={})

    # test allow missing
    mod.set_params(force_init=True, arg_params=arg_params_missing, aux_params={}, allow_missing=True)
    assertRaises(RuntimeError, mod.set_params,
                 force_init=True, arg_params=arg_params_missing,
                 aux_params={}, allow_missing=False)

    # test allow extra
    mod.set_params(force_init=True, arg_params=arg_params_extra, aux_params={}, allow_missing=True, allow_extra=True)
    assertRaises(ValueError, mod.set_params,
                 force_init=True, arg_params=arg_params_extra,
                 aux_params={}, allow_missing=True, allow_extra=False)

    # test allow missing + extra,
    assertRaises(RuntimeError, mod.set_params,
                 force_init=True, arg_params=arg_params_missing_extra,
                 aux_params={}, allow_missing=False, allow_extra=False)

    # test allow missing + extra, this will throw a runtime error
    assertRaises(ValueError, mod.set_params,
                 force_init=True, arg_params=arg_params_missing_extra,
                 aux_params={}, allow_missing=True, allow_extra=False)


@with_seed()
def test_monitor():
    # data iter
    data = mx.nd.array([[0.05, .10]]);
    label = mx.nd.array([[.01, 0.99]]);
    train_data = mx.io.NDArrayIter(data, label, batch_size=1)

    # symbols
    x = mx.symbol.Variable('data')
    x = mx.symbol.FullyConnected(name='fc_0', data=x, num_hidden=2)
    x = mx.symbol.Activation(name="act_0", data=x, act_type='sigmoid')
    x = mx.symbol.FullyConnected(name='fc_1', data=x, num_hidden=2)
    x = mx.symbol.Activation(name="act_1", data=x, act_type='sigmoid')
    x = mx.symbol.LinearRegressionOutput(data=x, name='softmax', grad_scale=2)

    # create monitor
    def mean_abs(x):
        sum_abs = mx.ndarray.sum(mx.ndarray.abs(x))
        return mx.ndarray.divide(sum_abs, reduce(lambda x, y: x * y, x.shape))
    mon = mx.mon.Monitor(1, stat_func=mean_abs, pattern='.*', sort=True)

    # create module
    mod = mx.mod.Module(x, context=[mx.cpu()]);
    mod.bind(train_data.provide_data, label_shapes=train_data.provide_label,
                    for_training=True)
    mod.install_monitor(mon)
    arg_params = {'fc_0_weight': mx.nd.array([[.15, .20], [.25, .30]]),
                  'fc_0_bias'  : mx.nd.array([.35, .35]),
                  'fc_1_weight': mx.nd.array([[.40, .45], [.50, .55]]),
                  'fc_1_bias'  : mx.nd.array([.60, .60])}
    mod.init_params(arg_params=arg_params)

    data_iter = iter(train_data)
    data_batch = next(data_iter)
    mon.tic()
    mod.forward_backward(data_batch)
    res = mon.toc()
    keys = ['act_0', 'act_1', 'data', 'fc_0', 'fc_1', 'softmax']
    mon_result_counts = [0, 0, 0, 0, 0, 0]
    assert(len(res) == 21)
    for n, k, v in res:
        for idx, key in enumerate(keys):
            if k.startswith(key):
                mon_result_counts[idx] += 1
                break
    assert(mon_result_counts == [2, 2, 1, 6, 6, 4])

@with_seed()
def test_executor_group():
    def get_rnn_sym(num_layers, num_words, num_hidden, num_embed, seq_len, sparse_embedding):
        stack = mx.rnn.SequentialRNNCell()
        for i in range(num_layers):
            stack.add(mx.rnn.LSTMCell(num_hidden=num_hidden, prefix='lstm_l%d_' % i))
        data = mx.sym.Variable('data')
        label = mx.sym.Variable('softmax_label')
        if sparse_embedding:
            embed_weight = mx.sym.Variable('embed_weight', stype='row_sparse')
            embed = mx.sym.contrib.SparseEmbedding(data=data, input_dim=num_words,
                                                   weight=embed_weight, output_dim=num_embed,
                                                   name='embed')
        else:
            embed = mx.sym.Embedding(data=data, input_dim=num_words,
                                     output_dim=num_embed, name='embed')

        stack.reset()
        outputs, states = stack.unroll(seq_len, inputs=embed, merge_outputs=True)

        pred = mx.sym.Reshape(outputs, shape=(-1, num_hidden))
        pred = mx.sym.FullyConnected(data=pred, num_hidden=num_words, name='pred')

        label = mx.sym.Reshape(label, shape=(-1,))
        pred = mx.sym.SoftmaxOutput(data=pred, label=label, name='softmax')
        return pred

    def test_shared_exec_group(exec_grp_shared, exec_grp_created, shared_arg_names=None,
                               extra_args=None, check_shared_grad=True):
        # Test shared data arrays
        for i in range(len(exec_grp_shared.execs)):
            # test same shared_data_arrays for two exec groups
            shared_data_array1 = exec_grp_shared.shared_data_arrays[i]
            shared_data_array2 = exec_grp_created.shared_data_arrays[i]
            if extra_args is not None:
                assert len(shared_data_array1) == len(extra_args),\
                    "exec_grp_shared.shared_data_arrays[%d] should have same number of args as extra_args"
            assert len(shared_data_array1) == len(shared_data_array2),\
                "length of shared_data_array of the shared executor group not equal to the created executor group"
            for k, v in shared_data_array1.items():
                if extra_args is not None:
                    assert k in extra_args, "arg %s is not in extra_args" % k
                assert k in shared_data_array2,\
                    "arg %s of the shared executor group not in the shared_data_array of the created executor group" % k
                assert mx.test_utils.same_array(v, shared_data_array2[k])

            for data_name, array in exec_grp_shared.shared_data_arrays[i].items():
                assert data_name in exec_grp_created.shared_data_arrays[i], \
                    "Shared input data '%s' is not in " \
                    "shared_data_arrays of created executor group." % (data_name)
                assert mx.test_utils.same_array(array, exec_grp_created.shared_data_arrays[i][data_name]), \
                    "Shared input data '%s' does not share memory." % (data_name)

            # Test shared argument arrays and gradient arrays
            exec_shared = exec_grp_shared.execs[i]
            exec_created = exec_grp_created.execs[i]
            if shared_arg_names is not None:
                # test shared arguments
                for arg_name in shared_arg_names:
                    assert arg_name in exec_created.arg_dict, \
                        "Shared argument '%s' is not in arg_dict of created executor group." % (arg_name)
                    assert mx.test_utils.same_array(exec_shared.arg_dict[arg_name], exec_created.arg_dict[arg_name]), \
                        "Shared argument '%s' does not share memory." % (arg_name)
                # test shared argument gradients
                if check_shared_grad:
                    for arg_name in shared_arg_names:
                        assert arg_name in exec_created.grad_dict, \
                            "Shared argument gradient '%s' is not in " \
                            "grad_dict of created executor group." % (arg_name)
                        assert mx.test_utils.same_array(exec_shared.grad_dict[arg_name], \
                                                        exec_created.grad_dict[arg_name]), \
                            "Shared argument gradient '%s' does not share memory." % (arg_name)

            for arg_name, grad in exec_grp_shared.grad_req.items():
                assert grad == exec_grp_created.grad_req[arg_name], \
                    "Gradient requirements for shared argument '%s' are inconsistent. " \
                    "Shared executor group requires '%s' while created executor group requires '%s'" \
                    %(arg_name, grad, exec_grp_created.grad_req[arg_name])

    def check_shared_exec_group(sparse_embedding):
        # generate an rnn sym with #layers=5
        sym = get_rnn_sym(num_layers=3, num_words=num_words, num_hidden=num_hidden,
                          num_embed=num_embed, seq_len=max_bucket_size,
                          sparse_embedding=sparse_embedding)
        arg_names1 = sym.list_arguments()
        input_names = [name[0] for name in data_shapes] + [name[0] for name in label_shapes]
        shared_arg_names = [name for name in arg_names1 if name not in input_names]
        exec_group1 = DataParallelExecutorGroup(symbol=sym, contexts=contexts,
                                                workload=workload, data_shapes=data_shapes,
                                                label_shapes=label_shapes, param_names=shared_arg_names,
                                                for_training=True, inputs_need_grad=False)

        # shared_data_arrays should only have input "data" and "softmax_label" arrays
        for i in range(len(contexts)):
            assert len(exec_group1.shared_data_arrays[i]) == len(input_names),\
                "exec_group1.shared_data_arrays[%d] should have the same number of names as in input_names" % i
            for name in input_names:
                assert name in exec_group1.shared_data_arrays[i],\
                    "arg %s should be in exec_group1.shared_data_arrays[%d]" % (name, i)

        # generate an rnn sym with #layers=5
        sym = get_rnn_sym(num_layers=5, num_words=num_words, num_hidden=num_hidden,
                          num_embed=num_embed, seq_len=max_bucket_size,
                          sparse_embedding=sparse_embedding)
        arg_names2 = sym.list_arguments()
        exec_group2 = DataParallelExecutorGroup(symbol=sym, contexts=contexts,
                                                workload=workload, data_shapes=data_shapes,
                                                label_shapes=label_shapes, param_names=shared_arg_names,
                                                for_training=True, inputs_need_grad=False,
                                                shared_group=exec_group1)
        extra_args = [name for name in arg_names2 if name not in shared_arg_names]
        check_shared_grad = not sparse_embedding
        test_shared_exec_group(exec_grp_shared=exec_group1, exec_grp_created=exec_group2,
                               shared_arg_names=shared_arg_names, extra_args=extra_args,
                               check_shared_grad=check_shared_grad)

    contexts = [mx.cpu(0), mx.cpu(1)]
    workload = [1] * len(contexts)
    batch_size = 32
    max_bucket_size = 80
    num_words = 1000
    num_hidden = 100
    num_embed = 200
    data_shapes = [('data', (batch_size, max_bucket_size))]
    label_shapes = [('softmax_label', (batch_size, max_bucket_size))]
    sparse_embedding_opt = [True, False]
    for opt in sparse_embedding_opt:
        check_shared_exec_group(opt)

@with_seed()
def test_factorization_machine_module():
    """ Test factorization machine model with sparse operators """
    # this unit test is to test the flow, training accuracy is tested in another test
    def check_factorization_machine_module(num_epochs=None):
        print("check_factorization_machine_module")

        def fm(factor_size, feature_dim, init):
            x = mx.symbol.Variable("data", stype='csr')
            v = mx.symbol.Variable("v", shape=(feature_dim, factor_size),
                                   init=init, stype='row_sparse')

            w1_weight = mx.symbol.var('w1_weight', shape=(feature_dim, 1),
                                      init=init, stype='row_sparse')
            w1_bias = mx.symbol.var('w1_bias', shape=(1))
            w1 = mx.symbol.broadcast_add(mx.symbol.dot(x, w1_weight), w1_bias)

            v_s = mx.symbol._internal._square_sum(data=v, axis=1, keepdims=True)
            x_s = mx.symbol.square(data=x)
            bd_sum = mx.sym.dot(x_s, v_s)

            w2 = mx.symbol.dot(x, v)
            w2_squared = 0.5 * mx.symbol.square(data=w2)

            w_all = mx.symbol.Concat(w1, w2_squared, dim=1)
            sum1 = mx.symbol.sum(data=w_all, axis=1, keepdims=True)
            sum2 = 0.5 * mx.symbol.negative(bd_sum)
            model = mx.sym.elemwise_add(sum1, sum2)

            y = mx.symbol.Variable("label")
            model = mx.symbol.LinearRegressionOutput(data=model, label=y)
            return model

        # model
        init = mx.initializer.Normal(sigma=0.01)
        factor_size = 4
        feature_dim = 10000
        model = fm(factor_size, feature_dim, init)

        # data iter
        num_batches = 5
        batch_size = 64
        num_samples = batch_size * num_batches
        # generate some random csr data
        csr_nd = rand_ndarray((num_samples, feature_dim), 'csr', 0.1)
        label = mx.nd.ones((num_samples,1))
        # the alternative is to use LibSVMIter
        train_iter = mx.io.NDArrayIter(data=csr_nd,
                                       label={'label':label},
                                       batch_size=batch_size,
                                       last_batch_handle='discard')
        # create module
        mod = mx.mod.Module(symbol=model, data_names=['data'], label_names=['label'])
        # allocate memory by given the input data and lable shapes
        mod.bind(data_shapes=train_iter.provide_data, label_shapes=train_iter.provide_label)
        # initialize parameters by uniform random numbers
        mod.init_params(initializer=init)

        # use Sparse SGD with learning rate 0.1 to train
        sgd = mx.optimizer.SGD(momentum=0.1, clip_gradient=5.0, learning_rate=0.01,
                               rescale_grad=1.0/batch_size)
        mod.init_optimizer(optimizer=sgd)
        if num_epochs is None:
            num_epochs = 50
        expected_accuracy = 0.02

	# use accuracy as the metric
        metric = mx.metric.create('MSE')
        # train 'num_epochs' epoch
        for epoch in range(num_epochs):
            train_iter.reset()
            metric.reset()
            for batch in train_iter:
                mod.forward(batch, is_train=True)       # compute predictions
                mod.update_metric(metric, batch.label)  # accumulate prediction accuracy
                mod.backward()                          # compute gradients
                mod.update()                            # update parameters
            print('Epoch %d, Training %s' % (epoch, metric.get()))
        if num_epochs > 1:
            assert(metric.get()[1] < expected_accuracy)

    check_factorization_machine_module()

@with_seed()
def test_module_initializer():
    def regression_model(m):
         x = mx.symbol.var("data", stype='csr')
         v = mx.symbol.var("v", shape=(m, 1), init=mx.init.Uniform(scale=.1),
                                stype='row_sparse')
         model = mx.symbol.dot(lhs=x, rhs=v)
         y = mx.symbol.Variable("label")
         model = mx.symbol.LinearRegressionOutput(data=model, label=y, name="out")
         return model

    n, m = 128, 100
    model = regression_model(m)

    data = mx.nd.zeros(shape=(n, m), stype='csr')
    label = mx.nd.zeros((n, 1))
    iterator = mx.io.NDArrayIter(data=data, label={'label':label},
                                 batch_size=n, last_batch_handle='discard')

    # create module
    mod = mx.mod.Module(symbol=model, data_names=['data'], label_names=['label'])
    mod.bind(data_shapes=iterator.provide_data, label_shapes=iterator.provide_label)
    mod.init_params()
    v = mod._arg_params['v']
    assert(v.stype == 'row_sparse')
    assert(np.sum(v.asnumpy()) != 0)

@with_seed()
def test_forward_reshape():
    num_class=10
    data1 = mx.sym.Variable('data1')
    data2 = mx.sym.Variable('data2')
    conv1 = mx.sym.Convolution(data=data1, kernel=(2, 2), num_filter=2, stride=(2, 2))
    conv2 = mx.sym.Convolution(data=data2, kernel=(3, 3), num_filter=3, stride=(1, 1))
    pooling1 = mx.sym.Pooling(data=conv1, kernel=(2, 2), stride=(1, 1), pool_type="avg")
    pooling2 = mx.sym.Pooling(data=conv2, kernel=(2, 2), stride=(1, 1), pool_type="max")
    flatten1 = mx.sym.flatten(data=pooling1)
    flatten2 = mx.sym.flatten(data=pooling2)
    sum = mx.sym.sum(data=flatten1, axis=1) + mx.sym.sum(data=flatten2, axis=1)
    fc = mx.sym.FullyConnected(data=sum, num_hidden=num_class)
    sym = mx.sym.SoftmaxOutput(data=fc, name='softmax')

    dshape1 = (10, 3, 64, 64)
    dshape2 = (10, 3, 32, 32)
    lshape = (10,)

    mod = mx.mod.Module(symbol=sym, data_names=['data1', 'data2'],
                        label_names=['softmax_label'])
    mod.bind(data_shapes=[('data1', dshape1), ('data2', dshape2)],
             label_shapes=[('softmax_label', lshape)])
    mod.init_params()
    mod.init_optimizer(optimizer_params={'learning_rate': 0.01})

    # Train with original data shapes
    data_batch = mx.io.DataBatch(data=[mx.nd.random.uniform(0, 9, dshape1),
                                       mx.nd.random.uniform(5, 15, dshape2)],
                                 label=[mx.nd.ones(lshape)])
    mod.forward(data_batch)
    assert mod.get_outputs()[0].shape == tuple([lshape[0], num_class])
    mod.backward()
    mod.update()

    # Train with different batch size
    dshape1 = (3, 3, 64, 64)
    dshape2 = (3, 3, 32, 32)
    lshape = (3,)
    data_batch = mx.io.DataBatch(data=[mx.nd.random.uniform(0, 9, dshape1),
                                       mx.nd.random.uniform(5, 15, dshape2)],
                                 label=[mx.nd.ones(lshape)])
    mod.forward(data_batch)
    assert mod.get_outputs()[0].shape == tuple([lshape[0], num_class])
    mod.backward()
    mod.update()

    dshape1 = (20, 3, 64, 64)
    dshape2 = (20, 3, 32, 32)
    lshape = (20,)
    data_batch = mx.io.DataBatch(data=[mx.nd.random.uniform(3, 5, dshape1),
                                       mx.nd.random.uniform(10, 25, dshape2)],
                                 label=[mx.nd.ones(lshape)])
    mod.forward(data_batch)
    assert mod.get_outputs()[0].shape == tuple([lshape[0], num_class])
    mod.backward()
    mod.update()

    #Train with both different batch size and data shapes
    dshape1 = (20, 3, 120, 120)
    dshape2 = (20, 3, 32, 64)
    lshape = (20,)
    data_batch = mx.io.DataBatch(data=[mx.nd.random.uniform(0, 9, dshape1),
                                       mx.nd.random.uniform(5, 15, dshape2)],
                                 label=[mx.nd.ones(lshape)])
    mod.forward(data_batch)
    assert mod.get_outputs()[0].shape == tuple([lshape[0], num_class])
    mod.backward()
    mod.update()

    dshape1 = (5, 3, 28, 40)
    dshape2 = (5, 3, 24, 16)
    lshape = (5,)
    data_batch = mx.io.DataBatch(data=[mx.nd.random.uniform(0, 9, dshape1),
                                       mx.nd.random.uniform(15, 25, dshape2)],
                                 label=[mx.nd.ones(lshape)])
    mod.forward(data_batch)
    assert mod.get_outputs()[0].shape == tuple([lshape[0], num_class])
    mod.backward()
    mod.update()

    #Test score
    dataset_shape1 = (30, 3, 30, 30)
    dataset_shape2 = (30, 3, 20, 40)
    labelset_shape = (30,)

    eval_dataiter = mx.io.NDArrayIter(data=[mx.nd.random.uniform(0, 9, dataset_shape1),
                                            mx.nd.random.uniform(15, 25, dataset_shape2)],
                                      label=[mx.nd.ones(labelset_shape)],
                                      batch_size=5)
    assert len(mod.score(eval_data=eval_dataiter, eval_metric='acc')) == 1

    #Test prediction
    dshape1 = (1, 3, 30, 30)
    dshape2 = (1, 3, 20, 40)
    dataset_shape1 = (10, 3, 30, 30)
    dataset_shape2 = (10, 3, 20, 40)

    pred_dataiter = mx.io.NDArrayIter(data=[mx.nd.random.uniform(0, 9, dataset_shape1),
                                            mx.nd.random.uniform(15, 25, dataset_shape2)])
    mod.bind(data_shapes=[('data1', dshape1), ('data2', dshape2)],
             for_training=False, force_rebind=True)
    assert mod.predict(pred_dataiter).shape == tuple([10, num_class])

@with_seed()
def test_forward_types():
    #Test forward with other data batch API
    Batch = namedtuple('Batch', ['data'])
    data = mx.sym.Variable('data')
    out = data * 2
    mod = mx.mod.Module(symbol=out, label_names=None)
    mod.bind(data_shapes=[('data', (1, 10))])
    mod.init_params()
    data1 = [mx.nd.ones((1, 10))]
    mod.forward(Batch(data1))
    assert mod.get_outputs()[0].shape == (1, 10)
    data2 = [mx.nd.ones((3, 5))]
    mod.forward(Batch(data2))
    assert mod.get_outputs()[0].shape == (3, 5)

    #Test forward with other NDArray and np.ndarray inputs
    data = mx.sym.Variable('data')
    out = data * 2
    mod = mx.mod.Module(symbol=out, label_names=None)
    mod.bind(data_shapes=[('data', (1, 10))])
    mod.init_params()
    data1 = mx.nd.ones((1, 10))
    assert mod.predict(data1).shape == (1, 10)
    data2 = np.ones((1, 10))
    assert mod.predict(data1).shape == (1, 10)


def test_reference_single_batch_during_fit():
    """
    When using C++-based iterators, it's important that only a single batch is referenced at a time. Because C++
    iterators are exposed to the Python code through a C API, there is no concept of reference counting. Hence,
    typically C++ iterators will deallocate a batch when next() is called on them. So, we need to make sure the Python
    code only references a single batch at a time, otherwise the Python code will attempt to access freed memory,
    resulting in either (a) garbage accuracy or (b) a segmentation fault.
    """
    current_batch_i = None

    class MockBatch(object):
        def __init__(self, i):
            self.i = i

        @property
        def label(self):
            global current_batch_i
            assert self.i == current_batch_i

    class MockTrainData(object):
        def __init__(self, batches):
            self._i = 0
            self._batches = batches
            self.provide_data = None
            self.provide_label = None
            self.reset = lambda: None

        def __iter__(self):
            self._i = 0
            return self

        def __next__(self):
            global current_batch_i

            if self._i < self._batches:
                current_batch_i = self._i
                self._i += 1
                return MockBatch(current_batch_i)
            raise StopIteration

        def next(self):
            return self.__next__()

    mod = mx.mod.BaseModule()

    def empty_fn(*args, **kwargs):
        pass
    mod.bind = empty_fn
    mod.init_params = empty_fn
    mod.init_optimizer = empty_fn
    mod.forward = empty_fn
    mod.backward = empty_fn
    mod.update = empty_fn
    mod.update_metric = empty_fn
    mod.get_params = lambda: (None, None)

    train_data = MockTrainData(batches=2)
    mod.fit(train_data, num_epoch=1)

@with_seed()
def test_bucket_module_grad_req():
    batch_size = 2
    def sym_gen(_):
        data = mx.symbol.Variable('data')
        weight = mx.symbol.Variable('a', shape=(1,), init=mx.init.One())
        sym = mx.sym.make_loss(mx.sym.broadcast_mul(data, weight))
        return sym, ('data',), None

    mod = mx.mod.BucketingModule(sym_gen=sym_gen, default_bucket_key=10)
    mod.bind(data_shapes=[['data', (batch_size, )]], for_training=True, grad_req='write')
    mod.init_params()

    mod.forward_backward(mx.io.DataBatch(data=[mx.nd.ones((batch_size,))],
                                         label=None,
                                         provide_data=[mx.io.DataDesc(name='data', shape=(batch_size, ), layout='N')],
                                         bucket_key=10))
    assert(mod._curr_module._exec_group.execs[0].grad_dict['a'].asscalar() == batch_size)

    mod.forward_backward(mx.io.DataBatch(data=[mx.nd.ones((batch_size,))],
                                         label=None,
                                         provide_data=[mx.io.DataDesc(name='data', shape=(batch_size, ), layout='N')],
                                         bucket_key=5))
    assert(mod._curr_module._exec_group.execs[0].grad_dict['a'].asscalar() == batch_size)

    mod = mx.mod.BucketingModule(sym_gen=sym_gen, default_bucket_key=10)
    mod.bind(data_shapes=[['data', (batch_size, )]], for_training=True, grad_req='add')
    mod.init_params()

    mod.forward_backward(mx.io.DataBatch(data=[mx.nd.ones((batch_size,))],
                                         label=None,
                                         provide_data=[mx.io.DataDesc(name='data', shape=(batch_size,), layout='N')],
                                         bucket_key=10))
    assert(mod._curr_module._exec_group.execs[0].grad_dict['a'].asscalar() == batch_size)

    mod.forward_backward(mx.io.DataBatch(data=[mx.nd.ones((batch_size,))],
                                         label=None,
                                         provide_data=[mx.io.DataDesc(name='data', shape=(batch_size,), layout='N')],
                                         bucket_key=5))
    assert mod._curr_module._grad_req == 'add'
    assert(mod._curr_module._exec_group.execs[0].grad_dict['a'].asscalar() == 2 * batch_size)


def test_module_update_no_pragram():
    # test module to do update on layers without params
    data_shape = (10, 10)
    data = mx.sym.Variable('data')
    out = mx.sym.Dropout(data, 0.5)
    mod = mx.mod.Module(out)
    mod.bind(data_shapes=[('data', data_shape)])
    mod.init_params()
    mod.init_optimizer()
    data_batch = mx.io.DataBatch([nd.ones(data_shape)])
    mod.forward_backward(data_batch)
    mod.update()
    assert(mod.get_outputs()[0].shape == data_shape)


def test_module_init_optimizer():
    def get_module_idx2name(mod):
        idx2name = {}
        idx2name.update(enumerate(mod._exec_group.param_names))
        return idx2name

    data = mx.sym.Variable('data')
    sym = mx.sym.FullyConnected(data, num_hidden=20, name='fc')
    batch_size = 8
    opt_params = {'learning_rate': 1, 'rescale_grad': 1.0 / batch_size}

    # Pass an optimizer str
    mod1 = mx.mod.Module(sym, ('data',), None, context=mx.cpu(0))
    mod1.bind(data_shapes=[('data', (batch_size, 20))])
    mod1.init_params()
    mod1.init_optimizer(optimizer='sgd', optimizer_params=opt_params)
    assert mod1._optimizer.idx2name == get_module_idx2name(mod1)

    # Pass an Optimizer object
    mod2 = mx.mod.Module(sym, ('data',), None, context=mx.cpu(0))
    mod2.bind(data_shapes=[('data', (batch_size, 20))])
    mod2.init_params()
    opt = mx.optimizer.SGD(**opt_params)
    mod2.init_optimizer(optimizer=opt)
    assert mod2._optimizer.idx2name == get_module_idx2name(mod2)


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