# 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 numpy as np
from common import with_seed, assertRaises
from mxnet.contrib.svrg_optimization.svrg_module import SVRGModule
from mxnet.test_utils import *
import pytest

def setup():
    train_data = np.random.randint(1, 5, [1000, 2])
    weights = np.array([1.0, 2.0])
    train_label = train_data.dot(weights)

    di = mx.io.NDArrayIter(train_data, train_label, batch_size=32, shuffle=True, label_name='lin_reg_label')
    X = mx.sym.Variable('data')
    Y = mx.symbol.Variable('lin_reg_label')
    fully_connected_layer = mx.sym.FullyConnected(data=X, name='fc1', num_hidden=1)
    lro = mx.sym.LinearRegressionOutput(data=fully_connected_layer, label=Y, name="lro")

    mod = SVRGModule(
        symbol=lro,
        data_names=['data'],
        label_names=['lin_reg_label'], update_freq=2)
    mod.bind(data_shapes=di.provide_data, label_shapes=di.provide_label)
    mod.init_params(initializer=mx.init.Uniform(0.01), allow_missing=False, force_init=False, allow_extra=False)

    return di, mod


def test_bind_module():
    _, mod = setup()
    assert mod.binded == True
    assert mod._mod_aux.binded == True


def test_module_init():
    _, mod = setup()
    assert mod._mod_aux is not None


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

    #shape of the data
    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 = SVRGModule(symbol=model, data_names=['data'], label_names=['label'], update_freq=2)
    mod.bind(data_shapes=iterator.provide_data, label_shapes=iterator.provide_label)
    mod.init_params()
    v = mod._arg_params['v']