import numpy as np
from mindspore import dataset as ds
from mindspore.common.initializer import Normal
from mindspore import nn
from mindspore.train import Model
from mindspore.train.callback import LossMonitor
from mindspore import context

context.set_context(mode=context.GRAPH_MODE, device_target="CPU")


def get_data(num, w=2.0, b=3.0):
    for i in range(num):
        x = np.random.uniform(-10.0, 10.0)
        noise = np.random.normal(0, 1)
        y = x * w + b + noise
        yield np.array([x]).astype(np.float32), np.array([y]).astype(np.float32)


def create_dataset(num_data, batch_size=16, repeat_size=1):
    input_data = ds.GeneratorDataset(list(get_data(num_data)), column_names=['data','label'])
    input_data = input_data.batch(batch_size)
    input_data = input_data.repeat(repeat_size)
    return input_data


class LinearNet(nn.Cell):
    def __init__(self):
        super(LinearNet, self).__init__()
        self.fc = nn.Dense(1, 1, Normal(0.02), Normal(0.02))

    def construct(self, x):
        x = self.fc(x)
        return x


if __name__ == "__main__":

    num_data = 1600
    batch_size = 16
    repeat_size = 1
    lr = 0.005
    momentum = 0.9
    
    net = LinearNet()
    net_loss = nn.loss.MSELoss()
    opt = nn.Momentum(net.trainable_params(), lr, momentum)
    model = Model(net, net_loss, opt)
    
    ds_train = create_dataset(num_data, batch_size=batch_size, repeat_size=repeat_size) 
    model.train(1, ds_train, callbacks=LossMonitor(), dataset_sink_mode=False)
    
    print(net.trainable_params()[0], "\n%s" % net.trainable_params()[1])